contextual-retrieval-for-user-memory¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/contextual-retrieval-for-user-memory
项目说明¶
实验 3-12：利用上下文感知检索增强用户记忆¶
将上下文感知检索技术应用于用户记忆的构建，是解决传统对话历史分块所面临的核心痛点，并迈向更高层次记忆能力的关键。本项目实现了一个双层记忆系统，结合了：
- 上下文感知检索（Contextual RAG）：对话历史的精准检索
- 高级 JSON 卡片（Advanced JSON Cards）：结构化的核心事实存储
🆕 最新更新¶
LLM-Based Memory Card Generation¶
- 自动提取：使用 LLM 从对话中智能提取结构化记忆卡片
- 完整结构：每张卡片包含 backstory、person、relationship 等必要字段
- 智能降级：当 LLM 不可用时自动降级到关键词提取
LLM Judge Integration¶
- 自动评估：集成 LLM Judge 对 Agent 回答进行自动评分
- 双路径支持：支持导入模块或直接 API 调用
- 详细反馈：提供 0-1 分数、通过/失败状态和评估理由
Enhanced Debugging¶
- 内存卡片可视化：评估时自动打印所有记忆卡片的完整 JSON
- 测试用例排序：按名称字母顺序显示测试用例
- 评估透明度：清晰显示 LLM Judge 使用状态
核心创新¶
1. 上下文增强的对话分块¶
传统的对话分块会丢失上下文信息。例如，一段孤立的对话片段"好的，就订这个吧"本身毫无信息量。只有知道上文是在讨论"从上海到西雅图的、价格为500美元的单程机票"，这段对话才有意义。
本系统在索引对话历史之前，增加了关键的"上下文生成"步骤： - 每个对话块都会调用 LLM 生成包含关键背景信息的前缀摘要 - 上下文包括时间、人物和意图等关键线索 - 极大提升了检索的准确性和相关性
2. 双层记忆结构¶
Advanced JSON Cards（常驻记忆） - 存储结构化的、总结性的核心事实 - 始终固定在 Agent 的上下文中 - 包含 backstory（信息来源）和 relationship（关联人员）等元数据 - 如："用户 Jessica 的护照将于2025年2月18日过期"
Contextual RAG（按需检索） - 提供对非结构化的原始对话细节的精准访问 - 快速找到具体讨论的完整上下文 - 作为决策的"证据"支持
3. LLM-Based Memory Extraction¶
系统现在能够从对话中智能提取结构化记忆卡片：
# 自动从对话生成记忆卡片
cards = indexer._generate_summary_cards(chunks, conversation_id)
# 生成的卡片示例：
{
"category": "financial",
"card_key": "bank_account_primary",
"backstory": "用户在开设账户时提供了银行信息",
"date_created": "2024-01-15 10:30:00",
"person": "John Smith (primary)",
"relationship": "primary account holder",
"bank_name": "Chase Bank",
"account_type": "checking",
"account_ending": "4567"
}
项目结构¶
contextual-retrieval-for-user-memory/
├── contextual_chunking.py # 上下文感知分块
├── advanced_memory_manager.py # 高级JSON卡片管理
├── contextual_indexer.py # 双层记忆索引器（含LLM提取）
├── contextual_agent.py # 结合双层记忆的Agent
├── contextual_evaluator.py # 评估框架（含LLM Judge）
├── contextual_compare.py # 🆕 离线对比脚本：上下文化 vs 原始块的召回（无需 API）
├── memory_qa_eval.json # 🆕 离线对比用的受控记忆问答对照集
├── main.py # 主入口（argparse，含 --mode compare 离线对比）
├── config.py # 配置管理
├── chunker.py # 基础分块器
├── tools.py # Agent工具
└── requirements.txt # 依赖项
安装与配置¶
1. 安装依赖¶
2. 配置环境变量¶
创建 .env 文件：
# LLM Provider Configuration
MOONSHOT_API_KEY=your_api_key_here
ARK_API_KEY=your_api_key_here
SILICONFLOW_API_KEY=your_api_key_here
OPENAI_API_KEY=your_api_key_here
# Default Provider
LLM_PROVIDER=kimi # Options: kimi, doubao, siliconflow, openai
# Model Settings
LLM_MODEL=kimi-k3 # 或其他模型
3. 启动检索管道服务¶
使用示例¶
离线对比：上下文化到底有没有用？（无需 API，推荐先跑这个）¶
本实验的核心论点是：在把对话记忆块送入嵌入/索引前，先为每块生成一段『上下文前缀』，能提升脱离上下文的孤立片段（如『好的，就订这个吧』）的召回。 --mode compare 提供一个完全离线、无需任何 API Key 或检索服务的受控对照实验来量化这一点。
它用同一份上下文，分别度量『不拼接（plain）』与『拼接后再索引（contextual）』两种方式的召回，变量只有『索引文本是否含上下文前缀』，因此结果直接反映上下文化本身的贡献。检索采用确定性的 BM25 词法检索（纯 Python、无第三方依赖）作为神经嵌入的离线代理；对照数据集见 memory_qa_eval.json（受控教学集，可用 --dataset 替换）。
# 打印对比指标表（Recall@1 / Recall@3 / MRR）
python main.py --mode compare
# 等价于直接运行独立脚本：
python contextual_compare.py
# 对单条查询做 plain vs contextual 的 Top-K 检索对比
python main.py --mode compare --query '我最后确认预订的那张机票是哪个航班？'
# 保存完整结果（含逐查询名次明细）到 JSON
python main.py --mode compare --output results/compare.json
实测输出（12 个记忆块、8 条查询的受控集）：
方法 Recall@1 Recall@3 MRR
--------------------------------------------------------------------
Plain（直接索引原始块） 0.625 1.000 0.792
Contextual（上下文化后索引） 0.750 1.000 0.875
--------------------------------------------------------------------
提升（Δ） +0.125 +0.000 +0.083
其中『我订的西雅图酒店确认了吗』这类查询，gold 记忆块（『可以，帮我订下来』）在 Plain 下排名第 3、上下文化后升到第 1——正是上下文前缀把孤立确认片段重新锚定回了『西雅图凯悦酒店』的情境。
说明：这是一个离线词法代理，用于在无 API 环境下清晰地演示上下文化的机制与方向性收益。生产管线中，上下文前缀由 LLM 逐块生成、并用神经嵌入 + 混合检索索引（见下方
--mode evaluate），需要配置 API Key。
端到端评估（需 API / 检索服务）¶
# 交互式界面（默认）
python main.py
# 评估特定分类（需 LLM 与检索管道服务）
python main.py --mode evaluate --category layer1
# 关闭上下文化做对照，并指定模型与输出
python main.py --mode evaluate --category layer1 --no-contextual --model gpt-5.6-luna --output results/plain_eval.json
完整命令行参数见 python main.py --help（含中文说明）。
交互式测试界面¶
运行 python main.py 进入交互式界面：
Main Menu:
1. 🚀 Demo Mode (Quick Start)
2. 📚 Load & Index Conversations
3. 🎴 Manage Memory Cards
4. 🔍 Test Query
5. 📊 Evaluate All Test Cases (by Category) [LLM Judge]
6. 🎯 Evaluate Specific Test Case [LLM Judge]
7. 📈 Show Statistics
8. ⚙️ Configure Settings
0. Exit
评估输出示例¶
============================================================
DEBUG: All Memory Cards in System
============================================================
[financial.bank_account_primary]
{
"backstory": "用户开设银行账户时提供的信息",
"date_created": "2024-06-12 14:30:00",
"person": "Michael James Robertson (primary)",
"relationship": "primary account holder",
"bank_name": "First National Bank",
"account_number": "4429853327",
"routing_number": "123006800"
}
Total Memory Cards: 5
============================================================
LLM Judge Evaluation Results
============================================================
Reward: 1.000/1.000
Passed: Yes
Reasoning: The agent correctly provided the checking account number...
============================================================
工作流程示例¶
当用户询问"为我一月的东京之行，还有什么要准备的吗？"时：
- 事实回顾：Agent 首先审视 Advanced JSON Cards 中的内容
- 发现"东京之行"信息（1月25日出发）
-
发现"护照信息"（2月18日过期）
-
关联与推理：通过对比核心事实
-
识别出机票日期与护照过期日期接近的风险
-
细节验证：启动 RAG 检索
- 搜索与"护照"和"东京机票"相关的对话片段
-
获取原始讨论的所有细节
-
主动服务：结合两种记忆
-
给出关键建议："您的护照即将过期，强烈建议您立即加急办理续签"
-
自动评估：LLM Judge 评估答案
- 评分：0.95/1.0
- 理由：正确识别风险并给出适当建议
参考资料¶
许可证¶
MIT License
OpenRouter 通用回退 / Universal OpenRouter fallback¶
This experiment now supports a universal OpenRouter fallback for its chat LLM.
- If the primary provider key (e.g.
MOONSHOT_API_KEY/KIMI_API_KEY/OPENAI_API_KEY/DOUBAO_API_KEY…) is present, behavior is unchanged. - Else if
OPENROUTER_API_KEYis set, the chat LLM is automatically routed through OpenRouter (https://openrouter.ai/api/v1). Model names are mapped automatically:gpt-*/o1-*→openai/…,claude-*→anthropic/claude-opus-4.8,kimi-*→moonshotai/kimi-k2.6, ids already containing/are kept as-is, and other provider-native ids (e.g.doubao-*) fall back toopenai/gpt-5.6-luna. SetOPENROUTER_MODELto force a specific OpenRouter model id. - Else a clear error lists the accepted keys.
Add OPENROUTER_API_KEY=... to your .env (see env.example) to enable it.
源代码¶
advanced_memory_manager.py¶
"""Advanced Memory Manager with JSON Cards for User Memory System
This module implements the Advanced JSON Cards approach from week2/user-memory,
which stores structured, summarized core facts about users.
"""
import json
import os
import uuid
import logging
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field, asdict
from pathlib import Path
logger = logging.getLogger(__name__)
@dataclass
class AdvancedMemoryCard:
"""
Represents an advanced memory card with complete metadata.
Each card MUST include:
- backstory: Context about when/why this information was learned
- date_created: Timestamp when created
- person: Who this relates to
- relationship: Role/relationship to primary user
- Additional fields based on information type
"""
category: str
card_key: str
backstory: str
date_created: str
person: str
relationship: str
data: Dict[str, Any] = field(default_factory=dict)
_metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary for storage"""
result = {
"backstory": self.backstory,
"date_created": self.date_created,
"person": self.person,
"relationship": self.relationship,
**self.data
}
if self._metadata:
result["_metadata"] = self._metadata
return result
@classmethod
def from_dict(cls, category: str, card_key: str, data: Dict[str, Any]) -> 'AdvancedMemoryCard':
"""Create from dictionary"""
metadata = data.pop("_metadata", {})
return cls(
category=category,
card_key=card_key,
backstory=data.get("backstory", ""),
date_created=data.get("date_created", ""),
person=data.get("person", ""),
relationship=data.get("relationship", ""),
data={k: v for k, v in data.items()
if k not in ["backstory", "date_created", "person", "relationship"]},
_metadata=metadata
)
class AdvancedMemoryManager:
"""
Advanced JSON Memory Manager for structured user information.
This manager handles the persistent, structured memory that stays
in the agent's context at all times (the "备忘录" or memo).
"""
def __init__(self, user_id: str, storage_dir: str = "memory_storage"):
"""
Initialize the advanced memory manager.
Args:
user_id: Unique user identifier
storage_dir: Directory for storing memory files
"""
self.user_id = user_id
self.storage_dir = Path(storage_dir)
self.storage_dir.mkdir(parents=True, exist_ok=True)
self.memory_file = self.storage_dir / f"{user_id}_advanced_memory.json"
self.categories: Dict[str, Dict[str, AdvancedMemoryCard]] = {}
self.load_memory()
logger.info(f"Initialized AdvancedMemoryManager for user {user_id}")
def load_memory(self):
"""Load memory cards from storage"""
if self.memory_file.exists():
try:
with open(self.memory_file, 'r', encoding='utf-8') as f:
data = json.load(f)
# Load categories and cards
for category, cards in data.get('categories', {}).items():
self.categories[category] = {}
for card_key, card_data in cards.items():
self.categories[category][card_key] = AdvancedMemoryCard.from_dict(
category, card_key, card_data
)
logger.info(f"Loaded {sum(len(cards) for cards in self.categories.values())} memory cards")
except Exception as e:
logger.error(f"Error loading memory: {e}")
self.categories = {}
else:
logger.info(f"No existing memory file for user {self.user_id}")
def save_memory(self):
"""Save memory cards to storage"""
try:
data = {
'user_id': self.user_id,
'type': 'advanced_json_cards',
'updated_at': datetime.now().isoformat(),
'categories': {}
}
# Convert cards to dict format
for category, cards in self.categories.items():
data['categories'][category] = {}
for card_key, card in cards.items():
data['categories'][category][card_key] = card.to_dict()
with open(self.memory_file, 'w', encoding='utf-8') as f:
json.dump(data, f, indent=2, ensure_ascii=False)
logger.info(f"Saved {sum(len(cards) for cards in self.categories.values())} memory cards")
except Exception as e:
logger.error(f"Error saving memory: {e}")
def add_card(self, card: AdvancedMemoryCard) -> str:
"""
Add a new memory card.
Args:
card: The memory card to add
Returns:
Memory ID in format category.card_key
"""
if card.category not in self.categories:
self.categories[card.category] = {}
# Add metadata
card._metadata = {
'created_at': datetime.now().isoformat(),
'updated_at': datetime.now().isoformat()
}
# Ensure required fields
if not card.backstory:
logger.warning("Memory card missing backstory")
if not card.date_created:
card.date_created = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
if not card.person:
card.person = "Primary User"
if not card.relationship:
card.relationship = "primary account holder"
self.categories[card.category][card.card_key] = card
self.save_memory()
memory_id = f"{card.category}.{card.card_key}"
logger.info(f"Added memory card: {memory_id}")
return memory_id
def update_card(self, category: str, card_key: str, updates: Dict[str, Any]) -> bool:
"""
Update an existing memory card.
Args:
category: Card category
card_key: Card key
updates: Fields to update
Returns:
True if successful
"""
if category not in self.categories or card_key not in self.categories[category]:
logger.warning(f"Card not found: {category}.{card_key}")
return False
card = self.categories[category][card_key]
# Update fields
for key, value in updates.items():
if key in ["backstory", "date_created", "person", "relationship"]:
setattr(card, key, value)
elif key != "_metadata":
card.data[key] = value
# Update metadata
if card._metadata:
card._metadata['updated_at'] = datetime.now().isoformat()
self.save_memory()
logger.info(f"Updated memory card: {category}.{card_key}")
return True
def delete_card(self, category: str, card_key: str) -> bool:
"""
Delete a memory card.
Args:
category: Card category
card_key: Card key
Returns:
True if successful
"""
if category in self.categories and card_key in self.categories[category]:
del self.categories[category][card_key]
# Clean up empty categories
if not self.categories[category]:
del self.categories[category]
self.save_memory()
logger.info(f"Deleted memory card: {category}.{card_key}")
return True
logger.warning(f"Card not found for deletion: {category}.{card_key}")
return False
def get_card(self, category: str, card_key: str) -> Optional[AdvancedMemoryCard]:
"""Get a specific memory card"""
if category in self.categories and card_key in self.categories[category]:
return self.categories[category][card_key]
return None
def search_cards(self, query: str) -> List[Tuple[str, AdvancedMemoryCard]]:
"""
Search memory cards by query.
Args:
query: Search query
Returns:
List of (memory_id, card) tuples
"""
query_lower = query.lower()
results = []
for category, cards in self.categories.items():
for card_key, card in cards.items():
memory_id = f"{category}.{card_key}"
# Search in all card fields
card_str = json.dumps(card.to_dict(), ensure_ascii=False).lower()
if (query_lower in category.lower() or
query_lower in card_key.lower() or
query_lower in card_str):
results.append((memory_id, card))
return results
def get_context_string(self, max_cards: Optional[int] = None) -> str:
"""
Get memory cards as formatted string for LLM context.
Args:
max_cards: Maximum number of cards to include (None for all)
Returns:
Formatted string of memory cards
"""
if not self.categories:
return "No advanced memory cards available."
lines = ["=== ADVANCED USER MEMORY CARDS ===\n"]
card_count = 0
for category, cards in self.categories.items():
lines.append(f"\n[Category: {category}]")
for card_key, card in cards.items():
if max_cards and card_count >= max_cards:
remaining = sum(len(c) for c in self.categories.values()) - card_count
lines.append(f"\n... and {remaining} more memory cards")
break
lines.append(f"\n Card '{card_key}':")
# Format card data nicely
card_dict = card.to_dict()
# Remove internal metadata from display
card_dict.pop('_metadata', None)
# Display key fields first
lines.append(f" - Backstory: {card.backstory}")
lines.append(f" - Person: {card.person}")
lines.append(f" - Relationship: {card.relationship}")
lines.append(f" - Date Created: {card.date_created}")
# Display other data fields
for key, value in card.data.items():
if isinstance(value, dict):
lines.append(f" - {key}:")
for k, v in value.items():
lines.append(f" {k}: {v}")
elif isinstance(value, list):
lines.append(f" - {key}: {', '.join(str(v) for v in value)}")
else:
lines.append(f" - {key}: {value}")
card_count += 1
if max_cards and card_count >= max_cards:
break
lines.append(f"\n\n[Total Memory Cards: {sum(len(cards) for cards in self.categories.values())}]")
return "\n".join(lines)
def summarize_for_conversation(self, conversation_id: str) -> Dict[str, Any]:
"""
Summarize memory cards for a specific conversation.
This creates a conversation-specific summary of relevant memory cards
that can be stored alongside the conversation chunks.
Args:
conversation_id: The conversation identifier
Returns:
Summary dictionary
"""
summary = {
"conversation_id": conversation_id,
"timestamp": datetime.now().isoformat(),
"total_cards": sum(len(cards) for cards in self.categories.values()),
"categories": {}
}
# Summarize by category
for category, cards in self.categories.items():
if cards:
summary["categories"][category] = {
"count": len(cards),
"keys": list(cards.keys()),
"sample": {}
}
# Include first 2 cards as samples
for i, (card_key, card) in enumerate(cards.items()):
if i >= 2:
break
summary["categories"][category]["sample"][card_key] = {
"person": card.person,
"backstory": card.backstory[:100] + "..." if len(card.backstory) > 100 else card.backstory
}
return summary
def clear_all_memories(self):
"""Clear all memory cards"""
self.categories = {}
self.save_memory()
logger.info(f"Cleared all memory cards for user {self.user_id}")
def get_statistics(self) -> Dict[str, Any]:
"""Get memory statistics"""
stats = {
"user_id": self.user_id,
"total_cards": sum(len(cards) for cards in self.categories.values()),
"categories": {}
}
for category, cards in self.categories.items():
stats["categories"][category] = {
"count": len(cards),
"cards": list(cards.keys())
}
# Calculate average backstory length
all_backstories = []
for cards in self.categories.values():
for card in cards.values():
if card.backstory:
all_backstories.append(len(card.backstory))
if all_backstories:
stats["avg_backstory_length"] = sum(all_backstories) / len(all_backstories)
else:
stats["avg_backstory_length"] = 0
return stats
def create_sample_cards() -> List[AdvancedMemoryCard]:
"""Create sample memory cards for testing"""
cards = [
AdvancedMemoryCard(
category="financial",
card_key="bank_account_primary",
backstory="User shared their banking details while setting up automatic bill payments",
date_created="2024-01-15 10:30:00",
person="John Smith (primary)",
relationship="primary account holder",
data={
"bank_name": "Chase Bank",
"account_type": "checking",
"account_ending": "4567",
"routing_number": "021000021",
"purpose": "primary checking for bills"
}
),
AdvancedMemoryCard(
category="medical",
card_key="doctor_dermatologist_sarah",
backstory="User needed to schedule a dermatology appointment for their daughter's skin condition",
date_created="2024-01-16 14:00:00",
person="Sarah Smith (daughter)",
relationship="family member",
data={
"doctor_name": "Dr. Emily Johnson",
"specialty": "Pediatric Dermatology",
"clinic": "Children's Health Center",
"phone": "555-0123",
"condition_treated": "eczema"
}
),
AdvancedMemoryCard(
category="travel",
card_key="passport_jessica",
backstory="User mentioned passport expiration while planning international travel",
date_created="2024-12-20 09:00:00",
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"passport_number": "XXXXX1234",
"expiration_date": "2025-02-18",
"issuing_country": "USA",
"needs_renewal": True
}
),
AdvancedMemoryCard(
category="travel",
card_key="tokyo_trip_january",
backstory="User booked a trip to Tokyo for a business conference",
date_created="2024-12-22 11:00:00",
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"destination": "Tokyo, Japan",
"departure_date": "2025-01-25",
"return_date": "2025-02-01",
"purpose": "business conference",
"hotel": "Hilton Tokyo",
"flight": "UA837"
}
)
]
return cards
agent.py¶
"""Agentic RAG Agent for User Memory Evaluation
This agent uses RAG-indexed conversation memories to answer questions
about user interactions, following the ReAct pattern.
"""
import json
import logging
from typing import List, Dict, Any, Optional, Generator
from dataclasses import dataclass, field
from datetime import datetime
from openai import OpenAI
from config import Config
from tools import MemoryTools, get_tool_definitions
from indexer import MemoryIndexer
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
@dataclass
class Message:
"""Represents a message in the conversation"""
role: str # "user", "assistant", "tool"
content: str
tool_calls: Optional[List[Dict[str, Any]]] = None
tool_call_id: Optional[str] = None
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
@dataclass
class AgentTrajectory:
"""Tracks the agent's reasoning and tool usage"""
test_id: str
question: str
iterations: List[Dict[str, Any]] = field(default_factory=list)
final_answer: Optional[str] = None
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
total_time: Optional[float] = None
success: bool = False
def to_dict(self) -> Dict[str, Any]:
return {
"test_id": self.test_id,
"question": self.question,
"iterations": self.iterations,
"final_answer": self.final_answer,
"tool_calls": self.tool_calls,
"total_time": self.total_time,
"success": self.success,
"total_iterations": len(self.iterations),
"total_tool_calls": len(self.tool_calls)
}
class UserMemoryRAGAgent:
"""Agent that uses RAG to answer questions about user conversation history"""
def __init__(self,
indexer: MemoryIndexer,
config: Optional[Config] = None):
"""
Initialize the agent
Args:
indexer: The memory indexer with loaded conversations
config: Configuration object
"""
self.config = config or Config.from_env()
self.indexer = indexer
self.memory_tools = MemoryTools(indexer)
# Initialize LLM client
self._init_llm_client()
# Tool definitions
self.tools = get_tool_definitions()
# Conversation history
self.conversation_history: List[Dict[str, Any]] = []
logger.info(f"Initialized UserMemoryRAGAgent with provider: {self.config.llm.provider}")
def _init_llm_client(self):
"""Initialize the LLM client based on provider"""
client_config, model = self.config.llm.get_client_config()
# Extract base_url if present
base_url = client_config.pop("base_url", None)
# Create OpenAI client
if base_url:
self.client = OpenAI(base_url=base_url, **client_config)
else:
self.client = OpenAI(**client_config)
self.model = model
logger.info(f"Using model: {self.model}")
def _get_system_prompt(self, test_id: str) -> str:
"""Generate the system prompt"""
return f"""You are an AI assistant with access to indexed conversation memories from user interactions.
Your task is to answer questions about these conversations accurately based ONLY on the information you can find in the indexed memories.
Current Test Case: {test_id}
## Important Guidelines:
1. **Memory Search Only**: You MUST only answer based on information found through the memory search tools. If the information is not available in the indexed conversations, clearly state that you cannot find it.
2. **Use Tools Effectively**:
- Use `search_memory` to find relevant information across all conversations
- Use `get_conversation_context` when you need more context around a search result
- Use `get_full_conversation` to review entire conversation histories when needed
3. **Multiple Searches**: Don't hesitate to perform multiple searches with different queries to find all relevant information. Different phrasings may yield different results.
4. **Citations Required**: Always mention which conversation or chunk you found the information in when providing answers.
5. **Be Thorough**: For complex questions, gather information from multiple chunks and conversations before formulating your answer.
6. **Handle Ambiguity**: If you find conflicting information or multiple possible answers, report all of them with their sources.
Remember: Your credibility depends on providing accurate, well-sourced information from the conversation memories only."""
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Execute a tool and return the result"""
try:
# Log tool call parameters to console
logger.info("="*80)
logger.info(f"TOOL CALL: {tool_name}")
logger.info(f"PARAMETERS: {json.dumps(arguments, indent=2, ensure_ascii=False)}")
logger.info("-"*80)
if tool_name == "search_memory":
query = arguments.get("query", "")
filter_test_id = arguments.get("filter_test_id")
result = self.memory_tools.search_memory(query, filter_test_id)
# Log result to console
result_dict = result.to_dict()
logger.info("TOOL RESULT:")
logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False))
logger.info("="*80)
return result_dict
elif tool_name == "get_conversation_context":
chunk_id = arguments.get("chunk_id", "")
context_size = arguments.get("context_size", 2)
result = self.memory_tools.get_conversation_context(chunk_id, context_size)
# Log result to console
result_dict = result.to_dict()
logger.info("TOOL RESULT:")
logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False))
logger.info("="*80)
return result_dict
elif tool_name == "get_full_conversation":
conversation_id = arguments.get("conversation_id", "")
test_id = arguments.get("test_id", "")
result = self.memory_tools.get_full_conversation(conversation_id, test_id)
# Log result to console
result_dict = result.to_dict()
logger.info("TOOL RESULT:")
logger.info(json.dumps(result_dict, indent=2, ensure_ascii=False))
logger.info("="*80)
return result_dict
else:
return {"status": "error", "error": f"Unknown tool: {tool_name}"}
except Exception as e:
logger.error(f"Tool execution error: {e}")
return {"status": "error", "error": str(e)}
def answer_question(self,
question: str,
test_id: str,
stream: bool = False) -> Dict[str, Any]:
"""
Answer a question about user conversation history using RAG
Args:
question: The question to answer
test_id: The test case ID for context
stream: Whether to stream the response
Returns:
Dictionary containing the answer and trajectory
"""
start_time = datetime.now()
trajectory = AgentTrajectory(test_id=test_id, question=question)
# Build initial messages
messages = [
{"role": "system", "content": self._get_system_prompt(test_id)},
{"role": "user", "content": question}
]
# Track iterations
iterations = 0
max_iterations = self.config.evaluation.max_iterations
# Process with ReAct loop
while iterations < max_iterations:
iterations += 1
iteration_data = {"iteration": iterations, "timestamp": datetime.now().isoformat()}
if self.config.agent.enable_reasoning:
logger.info(f"\n{'='*60}")
logger.info(f"Iteration {iterations}/{max_iterations}")
logger.info(f"{'='*60}")
try:
# Call LLM with tools
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=self.tools,
tool_choice="auto",
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
max_tokens=self.config.llm.max_tokens,
stream=False
)
message = response.choices[0].message
iteration_data["assistant_message"] = message.content or ""
# Log the LLM response content
if message.content:
logger.info("-"*60)
logger.info(f"LLM Response: {message.content}")
logger.info("-"*60)
# Add assistant message to history
assistant_msg = {"role": "assistant", "content": message.content or ""}
if message.tool_calls:
assistant_msg["tool_calls"] = [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
} for tc in message.tool_calls
]
iteration_data["tool_calls"] = []
messages.append(assistant_msg)
# Process tool calls if present
if message.tool_calls:
for tool_call in message.tool_calls:
tool_name = tool_call.function.name
try:
arguments = json.loads(tool_call.function.arguments)
except json.JSONDecodeError:
arguments = {}
# Execute tool
result = self._execute_tool(tool_name, arguments)
# Track tool call
tool_call_data = {
"tool": tool_name,
"arguments": arguments,
"result": result,
"timestamp": datetime.now().isoformat()
}
iteration_data["tool_calls"].append(tool_call_data)
trajectory.tool_calls.append(tool_call_data)
# Add tool result to messages
tool_message = {
"role": "tool",
"tool_call_id": tool_call.id,
"content": json.dumps(result, ensure_ascii=False)
}
messages.append(tool_message)
# Continue loop for next iteration
trajectory.iterations.append(iteration_data)
continue
else:
# No tool calls, we have final answer
trajectory.iterations.append(iteration_data)
trajectory.final_answer = message.content or ""
trajectory.success = True
# Calculate total time
end_time = datetime.now()
trajectory.total_time = (end_time - start_time).total_seconds()
# Return result
result = {
"answer": trajectory.final_answer,
"success": True,
"iterations": iterations,
"tool_calls": len(trajectory.tool_calls),
"trajectory": trajectory.to_dict() if self.config.evaluation.save_trajectories else None
}
if stream:
return self._stream_response(result)
else:
return result
except Exception as e:
logger.error(f"Error in iteration {iterations}: {e}")
trajectory.iterations.append({
"iteration": iterations,
"error": str(e)
})
# Continue to next iteration
continue
# Max iterations reached
logger.warning(f"Max iterations ({max_iterations}) reached")
# Calculate total time
end_time = datetime.now()
trajectory.total_time = (end_time - start_time).total_seconds()
trajectory.success = False
final_msg = "I was unable to find sufficient information to answer your question within the iteration limit. Please try rephrasing or breaking down your query."
return {
"answer": final_msg,
"success": False,
"iterations": iterations,
"tool_calls": len(trajectory.tool_calls),
"trajectory": trajectory.to_dict() if self.config.evaluation.save_trajectories else None
}
def _stream_response(self, result: Dict[str, Any]) -> Generator[str, None, None]:
"""Stream response content"""
# Stream the answer character by character
answer = result.get("answer", "")
for char in answer:
yield char
def clear_history(self):
"""Clear conversation history"""
self.conversation_history = []
logger.info("Conversation history cleared")
def save_trajectory(self, trajectory: AgentTrajectory, filepath: str):
"""
Save agent trajectory to file
Args:
trajectory: The trajectory to save
filepath: Path to save the trajectory
"""
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(trajectory.to_dict(), f, ensure_ascii=False, indent=2)
logger.info(f"Trajectory saved to {filepath}")
chunker.py¶
"""Conversation Chunker for User Memory RAG System
This module handles chunking of conversation histories into manageable segments
for indexing into the RAG database.
"""
import json
import logging
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field, asdict
from datetime import datetime
import hashlib
from config import ChunkingConfig, ChunkingStrategy
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class ConversationMessage:
"""Single message in a conversation"""
role: str # "user" or "assistant"
content: str
timestamp: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {k: v for k, v in asdict(self).items() if v is not None}
@dataclass
class ConversationChunk:
"""A chunk of conversation with metadata"""
chunk_id: str
conversation_id: str
test_id: str
chunk_index: int
start_round: int
end_round: int
messages: List[ConversationMessage]
metadata: Dict[str, Any] = field(default_factory=dict)
context_before: Optional[str] = None # Summary of previous context
context_after: Optional[str] = None # Preview of next context
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
def to_dict(self) -> Dict[str, Any]:
return {
"chunk_id": self.chunk_id,
"conversation_id": self.conversation_id,
"test_id": self.test_id,
"chunk_index": self.chunk_index,
"start_round": self.start_round,
"end_round": self.end_round,
"messages": [msg.to_dict() for msg in self.messages],
"metadata": self.metadata,
"context_before": self.context_before,
"context_after": self.context_after,
"created_at": self.created_at
}
def to_text(self) -> str:
"""Convert chunk to text format for indexing"""
lines = []
# Add metadata header
if self.metadata:
lines.append(f"[Conversation Metadata]")
for key, value in self.metadata.items():
lines.append(f" {key}: {value}")
lines.append("")
# Add context if available
if self.context_before:
lines.append(f"[Previous Context]\n{self.context_before}\n")
# Add messages
lines.append(f"[Conversation Rounds {self.start_round}-{self.end_round}]")
for msg in self.messages:
role_label = "Customer" if msg.role == "user" else "Representative"
lines.append(f"{role_label}: {msg.content}")
# Add next context preview if available
if self.context_after:
lines.append(f"\n[Next Context Preview]\n{self.context_after}")
return "\n".join(lines)
class ConversationChunker:
"""Chunks conversation histories into segments for RAG indexing"""
def __init__(self, config: Optional[ChunkingConfig] = None):
"""
Initialize the chunker
Args:
config: Chunking configuration
"""
self.config = config or ChunkingConfig()
logger.info(f"Initialized chunker with strategy: {self.config.strategy}")
logger.info(f"Rounds per chunk: {self.config.rounds_per_chunk}")
logger.info(f"Overlap rounds: {self.config.overlap_rounds}")
def chunk_conversation(self,
conversation_id: str,
test_id: str,
messages: List[Dict[str, Any]],
metadata: Optional[Dict[str, Any]] = None) -> List[ConversationChunk]:
"""
Chunk a single conversation into segments
Args:
conversation_id: Unique conversation identifier
test_id: Test case identifier
messages: List of messages in the conversation
metadata: Optional conversation metadata
Returns:
List of conversation chunks
"""
# Convert to ConversationMessage objects
conv_messages = []
for msg in messages:
conv_messages.append(ConversationMessage(
role=msg.get('role', 'user'),
content=msg.get('content', ''),
timestamp=msg.get('timestamp')
))
# Calculate rounds (1 round = 1 user message + 1 assistant response)
rounds = []
current_round = []
for msg in conv_messages:
current_round.append(msg)
if msg.role == "assistant" and len(current_round) >= 2:
rounds.append(current_round)
current_round = []
# Add remaining messages as incomplete round if any
if current_round:
rounds.append(current_round)
total_rounds = len(rounds)
logger.info(f"Processing conversation {conversation_id} with {total_rounds} rounds")
# Choose chunking strategy
if self.config.strategy == ChunkingStrategy.FIXED_ROUNDS:
chunks = self._chunk_fixed_rounds(
conversation_id, test_id, rounds, metadata
)
elif self.config.strategy == ChunkingStrategy.SEMANTIC:
# For now, fallback to fixed rounds
# Semantic chunking would require more sophisticated analysis
chunks = self._chunk_fixed_rounds(
conversation_id, test_id, rounds, metadata
)
else:
chunks = self._chunk_fixed_rounds(
conversation_id, test_id, rounds, metadata
)
logger.info(f"Created {len(chunks)} chunks for conversation {conversation_id}")
return chunks
def _chunk_fixed_rounds(self,
conversation_id: str,
test_id: str,
rounds: List[List[ConversationMessage]],
metadata: Optional[Dict[str, Any]] = None) -> List[ConversationChunk]:
"""
Chunk conversation using fixed number of rounds
Args:
conversation_id: Conversation identifier
test_id: Test case identifier
rounds: List of conversation rounds
metadata: Optional metadata
Returns:
List of chunks
"""
chunks = []
total_rounds = len(rounds)
# Calculate chunk boundaries with overlap
chunk_size = self.config.rounds_per_chunk
overlap = self.config.overlap_rounds
step = max(1, chunk_size - overlap)
chunk_index = 0
for start_idx in range(0, total_rounds, step):
end_idx = min(start_idx + chunk_size, total_rounds)
# Skip if chunk is too small (except for the last chunk)
if end_idx - start_idx < self.config.min_chunk_size and end_idx < total_rounds:
continue
# Flatten rounds into messages
chunk_messages = []
for round_idx in range(start_idx, end_idx):
chunk_messages.extend(rounds[round_idx])
# Generate chunk ID
chunk_content = f"{conversation_id}_{chunk_index}_{start_idx}_{end_idx}"
chunk_id = hashlib.md5(chunk_content.encode()).hexdigest()[:12]
# Create context summaries if enabled
context_before = None
context_after = None
if self.config.include_metadata:
# Add summary of previous context
if start_idx > 0:
prev_rounds = min(3, start_idx)
context_msgs = []
for i in range(max(0, start_idx - prev_rounds), start_idx):
for msg in rounds[i]:
if msg.role == "user":
context_msgs.append(f"User asked: {msg.content[:100]}...")
if context_msgs:
context_before = "Previous discussion: " + " | ".join(context_msgs[-2:])
# Add preview of next context
if end_idx < total_rounds:
next_rounds = min(2, total_rounds - end_idx)
context_msgs = []
for i in range(end_idx, min(end_idx + next_rounds, total_rounds)):
for msg in rounds[i]:
if msg.role == "user":
context_msgs.append(f"Next: {msg.content[:100]}...")
if context_msgs:
context_after = " | ".join(context_msgs[:2])
# Create chunk
chunk = ConversationChunk(
chunk_id=f"{test_id}_{conversation_id}_{chunk_id}",
conversation_id=conversation_id,
test_id=test_id,
chunk_index=chunk_index,
start_round=start_idx + 1, # 1-indexed for display
end_round=end_idx, # Inclusive
messages=chunk_messages,
metadata=metadata or {},
context_before=context_before,
context_after=context_after
)
chunks.append(chunk)
chunk_index += 1
# Stop if we've reached the end
if end_idx >= total_rounds:
break
return chunks
def chunk_test_case_conversations(self,
test_case: Dict[str, Any]) -> List[ConversationChunk]:
"""
Chunk all conversations in a test case
Args:
test_case: Test case containing conversation histories
Returns:
List of all chunks from all conversations
"""
all_chunks = []
test_id = test_case.get('test_id', 'unknown')
# Process each conversation history
for conv_history in test_case.get('conversation_histories', []):
conv_id = conv_history.get('conversation_id', '')
messages = conv_history.get('messages', [])
metadata = conv_history.get('metadata', {})
# Add test case information to metadata
metadata['test_id'] = test_id
metadata['test_title'] = test_case.get('title', '')
metadata['test_category'] = test_case.get('category', '')
# Chunk the conversation
chunks = self.chunk_conversation(
conversation_id=conv_id,
test_id=test_id,
messages=messages,
metadata=metadata
)
all_chunks.extend(chunks)
logger.info(f"Chunked test case {test_id}: {len(all_chunks)} total chunks")
return all_chunks
def save_chunks(self, chunks: List[ConversationChunk], filepath: str):
"""
Save chunks to a JSON file
Args:
chunks: List of conversation chunks
filepath: Path to save the chunks
"""
chunks_data = [chunk.to_dict() for chunk in chunks]
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(chunks_data, f, ensure_ascii=False, indent=2)
logger.info(f"Saved {len(chunks)} chunks to {filepath}")
def load_chunks(self, filepath: str) -> List[ConversationChunk]:
"""
Load chunks from a JSON file
Args:
filepath: Path to the chunks file
Returns:
List of conversation chunks
"""
with open(filepath, 'r', encoding='utf-8') as f:
chunks_data = json.load(f)
chunks = []
for chunk_data in chunks_data:
# Convert messages
messages = []
for msg_data in chunk_data.get('messages', []):
messages.append(ConversationMessage(**msg_data))
# Create chunk
chunk = ConversationChunk(
chunk_id=chunk_data['chunk_id'],
conversation_id=chunk_data['conversation_id'],
test_id=chunk_data['test_id'],
chunk_index=chunk_data['chunk_index'],
start_round=chunk_data['start_round'],
end_round=chunk_data['end_round'],
messages=messages,
metadata=chunk_data.get('metadata', {}),
context_before=chunk_data.get('context_before'),
context_after=chunk_data.get('context_after'),
created_at=chunk_data.get('created_at', '')
)
chunks.append(chunk)
logger.info(f"Loaded {len(chunks)} chunks from {filepath}")
return chunks
config.py¶
"""Configuration for Agentic RAG User Memory Evaluation System"""
import os
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List
from enum import Enum
from pathlib import Path
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
def _openrouter_model_id(model: Optional[str]) -> str:
"""Map a provider-native model name to an OpenRouter model id, used by the
universal OpenRouter fallback. An explicit OPENROUTER_MODEL env var wins."""
override = os.getenv("OPENROUTER_MODEL")
if override:
return override
m = (model or "").strip()
if not m:
return "openai/gpt-5.6-luna"
if "/" in m:
return m # already an OpenRouter-style id (e.g. openai/gpt-5.6-luna)
ml = m.lower()
if ml.startswith(("gpt-", "o1", "o3", "o4", "chatgpt")):
return "openai/" + m
if ml.startswith("claude-"):
return "anthropic/claude-opus-4.8"
if ml.startswith("kimi"):
# kimi-k3 is not on OpenRouter; moonshotai/kimi-k2.6 is the closest hosted id.
return "moonshotai/kimi-k2.6"
# Provider-native ids (kimi-*/doubao-*/qwen/deepseek-*) not hosted on
# OpenRouter under the same name -> a widely-available OpenAI chat model.
return "openai/gpt-5.6-luna"
class Provider(str, Enum):
"""Supported LLM providers"""
SILICONFLOW = "siliconflow"
DOUBAO = "doubao"
KIMI = "kimi"
MOONSHOT = "moonshot"
OPENROUTER = "openrouter"
OPENAI = "openai"
GROQ = "groq"
TOGETHER = "together"
DEEPSEEK = "deepseek"
class IndexMode(str, Enum):
"""Indexing modes for conversation chunks"""
DENSE = "dense" # Dense embedding only
SPARSE = "sparse" # Sparse embedding only (BM25)
HYBRID = "hybrid" # Both dense and sparse
class ChunkingStrategy(str, Enum):
"""Strategies for chunking conversations"""
FIXED_ROUNDS = "fixed_rounds" # Fixed number of rounds per chunk
SEMANTIC = "semantic" # Semantic boundaries
TIME_BASED = "time_based" # Based on timestamp gaps
@dataclass
class LLMConfig:
"""LLM configuration"""
provider: str = "kimi" # Default provider
model: Optional[str] = None # Will use provider defaults if not specified
api_key: Optional[str] = None # Will read from env if not provided
temperature: float = 0.7
max_tokens: int = 2048
stream: bool = True
# Provider-specific defaults
PROVIDER_DEFAULTS = {
"siliconflow": {
"model": "Qwen/Qwen3-235B-A22B-Thinking-2507",
"base_url": "https://api.siliconflow.cn/v1"
},
"doubao": {
"model": "doubao-seed-1-6-thinking-250715",
"base_url": "https://ark.cn-beijing.volces.com/api/v3"
},
"kimi": {
"model": "kimi-k3",
"base_url": "https://api.moonshot.cn/v1"
},
"moonshot": {
"model": "kimi-k3",
"base_url": "https://api.moonshot.cn/v1"
},
"openrouter": {
"model": "openai/gpt-5.6-luna",
"base_url": "https://openrouter.ai/api/v1"
},
"openai": {
"model": "gpt-5.6-luna",
"base_url": "https://api.openai.com/v1"
},
"groq": {
"model": "llama-3.3-70b-versatile",
"base_url": "https://api.groq.com/openai/v1"
},
"together": {
"model": "meta-llama/Llama-3.3-70B-Instruct-Turbo",
"base_url": "https://api.together.xyz"
},
"deepseek": {
"model": "deepseek-reasoner",
"base_url": "https://api.deepseek.com/v1"
}
}
def get_client_config(self) -> tuple[Dict[str, Any], str]:
"""Get OpenAI client configuration"""
provider = self.provider.lower()
defaults = self.PROVIDER_DEFAULTS.get(provider, {})
# Determine API key
api_key = self.api_key or os.getenv(f"{provider.upper()}_API_KEY")
if not api_key and provider == "moonshot":
api_key = os.getenv("KIMI_API_KEY") # Fallback for moonshot
# Determine model
model = self.model or defaults.get("model", "gpt-5.6-luna")
# Universal OpenRouter fallback: primary provider key absent but
# OPENROUTER_API_KEY present -> route through OpenRouter.
if not api_key and provider != "openrouter" and os.getenv("OPENROUTER_API_KEY"):
return {
"api_key": os.getenv("OPENROUTER_API_KEY"),
"base_url": "https://openrouter.ai/api/v1",
}, _openrouter_model_id(model)
# Build client config
client_config = {"api_key": api_key}
# Add base URL if needed
if base_url := defaults.get("base_url"):
client_config["base_url"] = base_url
return client_config, model
@dataclass
class ChunkingConfig:
"""Configuration for conversation chunking"""
strategy: ChunkingStrategy = ChunkingStrategy.FIXED_ROUNDS
rounds_per_chunk: int = 20 # Number of rounds per chunk for FIXED_ROUNDS
overlap_rounds: int = 2 # Number of overlapping rounds between chunks
include_metadata: bool = True # Include conversation metadata in chunks
min_chunk_size: int = 5 # Minimum number of rounds in a chunk
max_chunk_size: int = 50 # Maximum number of rounds in a chunk
@dataclass
class IndexConfig:
"""Configuration for RAG indexing"""
mode: IndexMode = IndexMode.HYBRID
embedding_model: str = "text-embedding-3-small" # OpenAI embedding model
embedding_dim: int = 1536 # Dimension of embeddings
index_path: str = "indexes/memory_index"
chunk_store_path: str = "data/chunk_store.json"
enable_contextual: bool = True # Add contextual information to chunks
contextual_window: int = 2 # Number of surrounding rounds for context
@dataclass
class EvaluationConfig:
"""Configuration for evaluation framework"""
test_cases_dir: str = "../../week2/user-memory-evaluation/test_cases"
results_dir: str = "results"
enable_verbose: bool = True
save_trajectories: bool = True
max_iterations: int = 10 # Max iterations for ReAct pattern
enable_caching: bool = True # Cache indexed conversations
use_llm_judge: bool = False # Use LLM to evaluate answers
@dataclass
class AgentConfig:
"""Agent behavior configuration"""
enable_reasoning: bool = True # Show reasoning steps
enable_citations: bool = True # Include citations in responses
max_search_results: int = 5 # Maximum search results to consider
confidence_threshold: float = 0.7 # Minimum confidence for answers
enable_multi_search: bool = True # Allow multiple searches per query
max_searches_per_query: int = 3 # Maximum searches allowed
verbose: bool = True # Enable verbose output
@dataclass
class Config:
"""Main configuration container"""
llm: LLMConfig = field(default_factory=LLMConfig)
chunking: ChunkingConfig = field(default_factory=ChunkingConfig)
index: IndexConfig = field(default_factory=IndexConfig)
evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
agent: AgentConfig = field(default_factory=AgentConfig)
@classmethod
def from_env(cls) -> "Config":
"""Create configuration from environment variables"""
config = cls()
# Override with environment variables
if provider := os.getenv("LLM_PROVIDER"):
config.llm.provider = provider
if model := os.getenv("LLM_MODEL"):
config.llm.model = model
if rounds := os.getenv("ROUNDS_PER_CHUNK"):
config.chunking.rounds_per_chunk = int(rounds)
if index_mode := os.getenv("INDEX_MODE"):
config.index.mode = IndexMode(index_mode)
return config
def save(self, path: str):
"""Save configuration to JSON file"""
import json
config_dict = {
"llm": {
"provider": self.llm.provider,
"model": self.llm.model,
"temperature": _reasoning_safe_temperature(self.llm.model, self.llm.temperature),
"max_tokens": self.llm.max_tokens,
"stream": self.llm.stream
},
"chunking": {
"strategy": self.chunking.strategy,
"rounds_per_chunk": self.chunking.rounds_per_chunk,
"overlap_rounds": self.chunking.overlap_rounds,
"include_metadata": self.chunking.include_metadata
},
"index": {
"mode": self.index.mode,
"embedding_model": self.index.embedding_model,
"enable_contextual": self.index.enable_contextual,
"contextual_window": self.index.contextual_window
},
"evaluation": {
"enable_verbose": self.evaluation.enable_verbose,
"save_trajectories": self.evaluation.save_trajectories,
"max_iterations": self.evaluation.max_iterations
},
"agent": {
"enable_reasoning": self.agent.enable_reasoning,
"enable_citations": self.agent.enable_citations,
"max_search_results": self.agent.max_search_results,
"confidence_threshold": self.agent.confidence_threshold
}
}
with open(path, 'w') as f:
json.dump(config_dict, f, indent=2)
@classmethod
def load(cls, path: str) -> "Config":
"""Load configuration from JSON file"""
import json
with open(path, 'r') as f:
config_dict = json.load(f)
config = cls()
# Update LLM config
if "llm" in config_dict:
for key, value in config_dict["llm"].items():
setattr(config.llm, key, value)
# Update other configs similarly
for section in ["chunking", "index", "evaluation", "agent"]:
if section in config_dict:
section_config = getattr(config, section)
for key, value in config_dict[section].items():
# Handle enums
if key == "strategy" and section == "chunking":
value = ChunkingStrategy(value)
elif key == "mode" and section == "index":
value = IndexMode(value)
setattr(section_config, key, value)
return config
contextual_agent.py¶
"""Contextual RAG Agent with Advanced Memory Cards
This agent combines:
1. Advanced Memory Cards (structured facts) - always in context
2. Contextual RAG for searching conversation history
"""
import json
import logging
import time
from typing import List, Dict, Any, Optional, Generator
from dataclasses import dataclass, field
from datetime import datetime
from openai import OpenAI
from config import Config
from contextual_indexer import ContextualMemoryIndexer
from advanced_memory_manager import AdvancedMemoryManager
from tools import MemoryTools, get_tool_definitions
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
@dataclass
class Message:
"""Represents a message in the conversation"""
role: str # "user", "assistant", "tool"
content: str
tool_calls: Optional[List[Dict[str, Any]]] = None
tool_call_id: Optional[str] = None
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
@dataclass
class AgentTrajectory:
"""Tracks the agent's reasoning and tool usage"""
test_id: str
question: str
iterations: List[Dict[str, Any]] = field(default_factory=list)
final_answer: Optional[str] = None
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
total_time: Optional[float] = None
success: bool = False
memory_cards_used: List[str] = field(default_factory=list)
chunks_retrieved: List[str] = field(default_factory=list)
def to_dict(self) -> Dict[str, Any]:
return {
"test_id": self.test_id,
"question": self.question,
"iterations": self.iterations,
"final_answer": self.final_answer,
"tool_calls": self.tool_calls,
"total_time": self.total_time,
"success": self.success,
"total_iterations": len(self.iterations),
"total_tool_calls": len(self.tool_calls),
"memory_cards_used": self.memory_cards_used,
"chunks_retrieved": self.chunks_retrieved
}
class ContextualUserMemoryAgent:
"""Agent with dual memory system: structured cards + contextual RAG"""
def __init__(self,
indexer: ContextualMemoryIndexer,
memory_manager: Optional[AdvancedMemoryManager] = None,
config: Optional[Config] = None):
"""
Initialize the contextual agent
Args:
indexer: The contextual memory indexer
memory_manager: Advanced memory manager (uses indexer's if not provided)
config: Configuration object
"""
self.config = config or Config.from_env()
self.indexer = indexer
self.memory_manager = memory_manager or indexer.memory_manager
self.memory_tools = MemoryTools(indexer) # Works with base indexer interface
# Set verbose flag
self.verbose = True # Always verbose for debugging
# Initialize LLM client
self._init_llm_client()
# Tool definitions - enhanced with contextual search
self.tools = self._get_enhanced_tool_definitions()
# Conversation history
self.conversation_history: List[Dict[str, Any]] = []
logger.info(f"Initialized ContextualUserMemoryAgent with dual memory system")
logger.info(f"Memory cards loaded: {sum(len(cards) for cards in self.memory_manager.categories.values())}")
def _init_llm_client(self):
"""Initialize the LLM client based on provider"""
client_config, model = self.config.llm.get_client_config()
# Extract base_url if present
base_url = client_config.pop("base_url", None)
# Create OpenAI client
if base_url:
self.client = OpenAI(base_url=base_url, **client_config)
else:
self.client = OpenAI(**client_config)
self.model = model
logger.info(f"Using model: {self.model}")
def _get_enhanced_tool_definitions(self) -> List[Dict[str, Any]]:
"""Get enhanced tool definitions for contextual search"""
return [
{
"type": "function",
"function": {
"name": "search_conversation_history",
"description": "Search through indexed conversation history with contextual understanding. Returns conversation chunks with their contextual descriptions.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query to find relevant conversations"
},
"top_k": {
"type": "integer",
"description": "Number of results to return",
"default": 3
}
},
"required": ["query"]
}
}
}
]
def _build_system_prompt(self) -> str:
"""Build enhanced system prompt with memory cards"""
# Get memory cards context (always included in prompt)
memory_context = self.memory_manager.get_context_string(max_cards=20)
prompt = f"""You are an intelligent assistant with access to comprehensive user memory:
{memory_context}
=== YOUR CAPABILITIES ===
1. **Memory Cards** (shown above): Pre-loaded structured facts about the user
- These are verified, persistent facts with backstories
- Each card shows WHO it's about and WHY we know this
- Always check these FIRST before searching
2. **Searchable Conversation History**: Use the search tool to find specific details
- Conversations are chunked with contextual descriptions
- Search when you need evidence or additional details
- Each chunk includes context about what's being discussed
=== PROACTIVE SERVICE GUIDELINES ===
You should provide proactive service by:
1. **Anticipating Needs**: Look beyond the immediate question to identify related concerns
2. **Risk Detection**: Identify potential issues before they become problems
- Check dates for expirations (passports, licenses, cards, insurances)
- Notice scheduling conflicts or tight timelines
- Flag missing preparations or requirements
3. **Comprehensive Assistance**: Connect different pieces of information
- If user asks about travel, check passport, visa, insurance status
- If discussing finances, consider upcoming payments or deadlines
- For medical topics, recall relevant history and upcoming appointments
4. **Helpful Suggestions**: Offer actionable recommendations
- Prioritize urgent matters (e.g., for time-sensitive issues)
- Suggest next steps even if not explicitly requested
- Remind about related tasks that might be overlooked
=== OPERATIONAL APPROACH ===
1. **Direct Answer First**: Address the user's immediate question clearly
2. **Then Proactive Service**: After answering, consider what else might be relevant
3. **Cross-Reference Information**: Actively connect related memory cards and conversations
4. **Cite Sources**: "According to memory card X..." or "Based on conversation Y..."
5. **Handle Conflicts**: Prefer more recent or more specific information
6. **Identify People**: Be specific about WHO information relates to
=== YOUR MISSION ===
Not just to answer questions, but to be a thoughtful assistant who:
- Notices what the user might have forgotten
- Warns about potential issues before they arise
- Provides comprehensive support beyond what's asked
- Acts as a reliable memory partner who cares about the user's wellbeing
When answering, always consider: "What else should the user know about this topic?"
Remember: Good service answers the question. Great service anticipates what comes next."""
return prompt
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
"""Execute a tool and return results"""
try:
if tool_name == "search_conversation_history":
query = arguments.get("query", "")
top_k = arguments.get("top_k", 3)
try:
# Use contextual search
search_results = self.indexer.search_with_context(
query=query,
top_k=top_k,
include_memory_cards=False # Cards are already in context
)
except Exception as search_error:
logger.error(f"Search failed: {search_error}")
return {
"status": "error",
"message": f"Search failed: {str(search_error)}",
"results": []
}
chunk_results = search_results.get("chunk_results", [])
if not chunk_results:
return {
"status": "success",
"message": "No relevant conversations found",
"results": []
}
# Format results with context - NO TRUNCATION
formatted_results = []
for result in chunk_results:
formatted_results.append({
"chunk_id": result.get("chunk_id"),
"context": result.get("context", ""), # Contextual description
"conversation": result.get("conversation_id"),
"rounds": result.get("rounds"),
"content": result.get("text", "") # Full content, no truncation
})
return {
"status": "success",
"results": formatted_results,
"total": len(formatted_results)
}
else:
return {
"status": "error",
"message": f"Unknown tool: {tool_name}"
}
except Exception as e:
logger.error(f"Tool execution error: {e}")
return {
"status": "error",
"message": str(e)
}
def answer_question(self,
question: str,
test_id: str = "unknown",
max_iterations: int = 10,
stream: bool = False) -> AgentTrajectory:
"""
Answer a question using dual memory system
Args:
question: The question to answer
test_id: Test case ID for tracking
max_iterations: Maximum reasoning iterations
stream: Whether to stream the response
Returns:
AgentTrajectory with the answer and reasoning steps
"""
trajectory = AgentTrajectory(test_id=test_id, question=question)
start_time = time.time()
# Reset conversation for new question
self.conversation_history = []
# Build initial messages with system prompt
messages = [
{"role": "system", "content": self._build_system_prompt()},
{"role": "user", "content": question}
]
# Track which memory cards might be relevant
relevant_cards = self._find_relevant_memory_cards(question)
trajectory.memory_cards_used = relevant_cards
iteration = 0
while iteration < max_iterations:
iteration += 1
iteration_data = {
"iteration": iteration,
"timestamp": datetime.now().isoformat(),
"messages_count": len(messages)
}
try:
# Generate response
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
tools=self.tools,
tool_choice="auto",
temperature=_reasoning_safe_temperature(self.model, 0.3),
max_tokens=2048,
stream=stream
)
if stream:
# Handle streaming response
assistant_content = ""
tool_calls = []
for chunk in response:
if chunk.choices[0].delta.content:
content = chunk.choices[0].delta.content
assistant_content += content
if self.verbose:
print(content, end='', flush=True)
# Handle tool calls in stream
if chunk.choices[0].delta.tool_calls:
for tc in chunk.choices[0].delta.tool_calls:
# Build tool calls from stream
pass # Simplified for brevity
# Process complete response
assistant_message = {
"role": "assistant",
"content": assistant_content if assistant_content else None
}
else:
# Non-streaming response
choice = response.choices[0]
assistant_message = {
"role": "assistant",
"content": choice.message.content
}
# Check for tool calls
if choice.message.tool_calls:
assistant_message["tool_calls"] = [
tc.model_dump() for tc in choice.message.tool_calls
]
messages.append(assistant_message)
iteration_data["response"] = assistant_message
# Handle tool calls
if assistant_message.get("tool_calls"):
if self.verbose:
print(f"\n{'='*80}")
print(f"🤖 LLM MADE {len(assistant_message['tool_calls'])} TOOL CALL(S)")
print(f"{'='*80}")
print("Tool calls:")
for tc in assistant_message["tool_calls"]:
print(f" - {tc['function']['name']}")
print()
for tool_call in assistant_message["tool_calls"]:
tool_name = tool_call["function"]["name"]
tool_args = json.loads(tool_call["function"]["arguments"])
# Execute tool
tool_result = self._execute_tool(tool_name, tool_args)
# Track tool usage
trajectory.tool_calls.append({
"iteration": iteration,
"tool": tool_name,
"arguments": tool_args,
"result": tool_result
})
# Track retrieved chunks
if tool_name == "search_conversation_history" and tool_result.get("results"):
for result in tool_result["results"]:
trajectory.chunks_retrieved.append(result.get("chunk_id", ""))
# Add tool result to messages
messages.append({
"role": "tool",
"tool_call_id": tool_call["id"],
"content": json.dumps(tool_result)
})
if self.verbose:
print(f"\n{'='*80}")
print(f"🔧 TOOL CALL: {tool_name}")
print(f"{'='*80}")
print(f"📥 Arguments:")
print(json.dumps(tool_args, indent=2, ensure_ascii=False))
print(f"\n📤 Result (FULL - NO TRUNCATION):")
print(json.dumps(tool_result, indent=2, ensure_ascii=False))
print(f"{'='*80}\n")
else:
# No tool calls means we have the final answer
trajectory.final_answer = assistant_message.get("content", "")
trajectory.success = True
break
trajectory.iterations.append(iteration_data)
except Exception as e:
logger.error(f"Error in iteration {iteration}: {e}")
iteration_data["error"] = str(e)
trajectory.iterations.append(iteration_data)
break
# Record timing
trajectory.total_time = time.time() - start_time
# Store conversation history
self.conversation_history = messages
if self.verbose:
print(f"\n{'='*80}")
print(f"✅ EVALUATION COMPLETE")
print(f"{'='*80}")
print(f"Iterations: {iteration}")
print(f"Total Time: {trajectory.total_time:.2f}s")
print(f"Memory Cards Used: {len(trajectory.memory_cards_used)}")
if trajectory.memory_cards_used:
print(f" Cards: {trajectory.memory_cards_used}")
print(f"Chunks Retrieved: {len(trajectory.chunks_retrieved)}")
if trajectory.chunks_retrieved:
print(f" Chunks: {trajectory.chunks_retrieved}")
print(f"\n📝 FINAL ANSWER:")
print(trajectory.final_answer or "No answer generated")
print(f"{'='*80}\n")
return trajectory
def _find_relevant_memory_cards(self, question: str) -> List[str]:
"""Find which memory cards might be relevant to the question"""
relevant = []
question_lower = question.lower()
# Simple keyword matching (could be enhanced with embeddings)
for category, cards in self.memory_manager.categories.items():
for card_key, card in cards.items():
card_str = json.dumps(card.to_dict()).lower()
# Check if any question keywords appear in the card
if any(word in card_str for word in question_lower.split() if len(word) > 3):
relevant.append(f"{category}.{card_key}")
return relevant
def reset(self):
"""Reset the agent state"""
self.conversation_history = []
logger.info("Agent state reset")
def get_statistics(self) -> Dict[str, Any]:
"""Get agent statistics"""
return {
"memory_cards": sum(len(cards) for cards in self.memory_manager.categories.values()),
"indexed_chunks": len(self.indexer.contextual_chunks),
"conversation_history_length": len(self.conversation_history),
"indexer_stats": self.indexer.get_statistics()
}
contextual_chunking.py¶
"""Contextual Chunking for User Memory Conversations
This module implements Anthropic's Contextual Retrieval approach specifically
for conversation histories. Each conversation chunk gets contextualized with
information about its position and meaning within the full conversation.
"""
import json
import hashlib
import logging
import time
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from openai import OpenAI
from config import ChunkingConfig, LLMConfig
from chunker import ConversationChunk, ConversationMessage
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
@dataclass
class ContextualConversationChunk:
"""Enhanced conversation chunk with contextual information"""
chunk_id: str
conversation_id: str
test_id: str
chunk_index: int
start_round: int
end_round: int
messages: List[ConversationMessage]
original_text: str # Original conversation text
context: str # Generated contextual description
contextualized_text: str # Context + original text
context_tokens: int = 0 # Track token usage
generation_time: float = 0.0 # Track generation time
metadata: Dict[str, Any] = field(default_factory=dict)
context_before: Optional[str] = None
context_after: Optional[str] = None
created_at: str = field(default_factory=lambda: datetime.now().isoformat())
def to_dict(self) -> Dict[str, Any]:
return {
"chunk_id": self.chunk_id,
"conversation_id": self.conversation_id,
"test_id": self.test_id,
"chunk_index": self.chunk_index,
"start_round": self.start_round,
"end_round": self.end_round,
"messages": [msg.to_dict() for msg in self.messages],
"original_text": self.original_text,
"context": self.context,
"contextualized_text": self.contextualized_text,
"context_tokens": self.context_tokens,
"generation_time": self.generation_time,
"metadata": self.metadata,
"context_before": self.context_before,
"context_after": self.context_after,
"created_at": self.created_at
}
@classmethod
def from_basic_chunk(cls, chunk: ConversationChunk, context: str = "",
contextualized_text: str = "", context_tokens: int = 0,
generation_time: float = 0.0) -> 'ContextualConversationChunk':
"""Create from a basic ConversationChunk"""
original_text = chunk.to_text()
if not contextualized_text:
contextualized_text = f"{context}\n\n{original_text}" if context else original_text
return cls(
chunk_id=chunk.chunk_id,
conversation_id=chunk.conversation_id,
test_id=chunk.test_id,
chunk_index=chunk.chunk_index,
start_round=chunk.start_round,
end_round=chunk.end_round,
messages=chunk.messages,
original_text=original_text,
context=context,
contextualized_text=contextualized_text,
context_tokens=context_tokens,
generation_time=generation_time,
metadata=chunk.metadata,
context_before=chunk.context_before,
context_after=chunk.context_after,
created_at=chunk.created_at
)
class ContextualConversationChunker:
"""
Implements contextual chunking specifically for conversation histories.
Key Educational Points:
1. Conversation Context: Each chunk gets context about what was discussed before and after
2. Topic Continuity: Context helps maintain topic understanding across chunk boundaries
3. Temporal Information: Context includes temporal relationships between conversations
4. Semantic Enrichment: Context adds semantic meaning to isolated conversation fragments
"""
def __init__(self,
chunking_config: Optional[ChunkingConfig] = None,
llm_config: Optional[LLMConfig] = None,
use_contextual: bool = True):
"""
Initialize the contextual conversation chunker.
Args:
chunking_config: Configuration for chunking parameters
llm_config: LLM configuration for context generation
use_contextual: Whether to generate contextual chunks (for comparison)
"""
self.chunking_config = chunking_config or ChunkingConfig()
self.llm_config = llm_config or LLMConfig()
self.use_contextual = use_contextual
# Initialize LLM client for context generation
if self.use_contextual:
self._init_llm_client()
# Statistics tracking
self.stats = {
"total_chunks": 0,
"contextual_chunks": 0,
"total_context_tokens": 0,
"total_generation_time": 0.0,
"cache_hits": 0,
"cache_misses": 0
}
# Context cache to avoid regenerating for similar chunks
self.context_cache = {}
logger.info(f"Initialized ContextualConversationChunker (contextual={use_contextual})")
def _init_llm_client(self):
"""Initialize LLM client for context generation"""
client_config, model = self.llm_config.get_client_config()
base_url = client_config.pop("base_url", None)
if base_url:
self.client = OpenAI(base_url=base_url, **client_config)
else:
self.client = OpenAI(**client_config)
self.model = model
logger.info(f"Using {self.llm_config.provider} ({self.model}) for context generation")
def contextualize_chunks(self,
chunks: List[ConversationChunk],
full_conversation_history: Optional[str] = None,
on_chunk_ready: Optional[callable] = None) -> List[ContextualConversationChunk]:
"""
Add contextual information to conversation chunks.
Args:
chunks: List of basic conversation chunks
full_conversation_history: Optional full conversation text for better context
on_chunk_ready: Callback for when each chunk is ready
Returns:
List of contextualized conversation chunks
"""
if not self.use_contextual:
# Return non-contextual chunks for comparison
return [
ContextualConversationChunk.from_basic_chunk(chunk)
for chunk in chunks
]
contextual_chunks = []
logger.info(f"Contextualizing {len(chunks)} conversation chunks")
# Build full conversation if not provided
if full_conversation_history is None:
full_conversation_history = self._build_full_conversation(chunks)
for i, chunk in enumerate(chunks):
logger.info(f"Generating context for chunk {i+1}/{len(chunks)}")
# Generate context for this chunk
chunk_hash = hashlib.md5(chunk.to_text().encode()).hexdigest()
if chunk_hash in self.context_cache:
context = self.context_cache[chunk_hash]
generation_time = 0.0
context_tokens = 0
self.stats["cache_hits"] += 1
logger.debug(f"Cache hit for chunk {chunk.chunk_id}")
else:
# Generate new context using Anthropic's approach
context, context_tokens, generation_time = self._generate_chunk_context(
chunk,
full_conversation_history,
)
# Cache the context
self.context_cache[chunk_hash] = context
self.stats["cache_misses"] += 1
self.stats["total_context_tokens"] += context_tokens
self.stats["total_generation_time"] += generation_time
# Create contextual chunk
contextual_chunk = ContextualConversationChunk.from_basic_chunk(
chunk,
context=context,
context_tokens=context_tokens,
generation_time=generation_time
)
contextual_chunks.append(contextual_chunk)
# Log the full contextualized chunk
logger.info(f"\n{'='*80}")
logger.info(f"📝 CONTEXTUAL CHUNK {i}/{len(chunks)} CREATED")
logger.info(f"{'='*80}")
logger.info(f"Chunk ID: {contextual_chunk.chunk_id}")
logger.info(f"Rounds: {contextual_chunk.start_round}-{contextual_chunk.end_round}")
logger.info(f"Context Generated:\n{context}")
logger.info(f"Original Text Preview (first 500 chars):\n{contextual_chunk.original_text[:500]}...")
logger.info(f"Context Tokens: {context_tokens}")
logger.info(f"Generation Time: {generation_time:.2f}s")
logger.info(f"{'='*80}\n")
# Call callback if provided
if on_chunk_ready:
try:
on_chunk_ready(contextual_chunk)
except Exception as e:
logger.error(f"Error in on_chunk_ready callback: {e}")
# Log progress
if (i + 1) % 5 == 0:
avg_time = self.stats["total_generation_time"] / (i + 1) if i > 0 else 0
logger.info(f"Progress: {i+1}/{len(chunks)} chunks, avg time: {avg_time:.2f}s")
# Update statistics
self.stats["total_chunks"] += len(contextual_chunks)
self.stats["contextual_chunks"] += len(contextual_chunks)
logger.info(f"Contextualization complete. Statistics: {json.dumps(self.stats, indent=2)}")
return contextual_chunks
def _build_full_conversation(self, chunks: List[ConversationChunk]) -> str:
"""Build full conversation text from chunks"""
all_messages = []
for chunk in chunks:
all_messages.extend([chunk.to_text()])
return "\n\n---\n\n".join(all_messages)
def _generate_chunk_context(self,
chunk: ConversationChunk,
full_conversation: str,
) -> Tuple[str, int, float]:
"""
Generate contextual description for a conversation chunk.
This follows Anthropic's Contextual Retrieval approach adapted for conversations:
1. Provide the full conversation history
2. Show the specific chunk
3. Ask for concise context about what's being discussed
Returns:
Tuple of (context, token_count, generation_time)
"""
start_time = time.time()
try:
# Build context about surrounding conversation
chunk_text = chunk.to_text()
# Create prompt following Anthropic's template with conversation-specific adaptations
prompt = f"""<full_conversation>
{full_conversation}
</full_conversation>
Here is a specific conversation chunk we want to situate within the whole conversation history:
<chunk>
Conversation rounds {chunk.start_round} to {chunk.end_round}:
{chunk_text}
</chunk>
Additional context:
- This is chunk {chunk.chunk_index + 1} from conversation {chunk.conversation_id}
Please provide a brief context (2-3 sentences) that describes:
1. What main topics or issues are being discussed in this chunk
2. Any key decisions, facts, or preferences mentioned
3. How this relates to the overall conversation flow
Write the context in a way that would help someone quickly understand what this conversation segment is about when searching through conversation history. Use the same language as the conversation."""
# Generate context
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3), # Low temperature for consistency
max_tokens=150 # Slightly more tokens for conversation context
)
context = response.choices[0].message.content.strip()
# Estimate token count (rough approximation)
token_count = len(prompt.split()) + len(context.split())
generation_time = time.time() - start_time
# Add structured prefix to context
context_prefix = f"[Conversation {chunk.conversation_id}, Rounds {chunk.start_round}-{chunk.end_round}] "
context = context_prefix + context
logger.debug(f"Generated context in {generation_time:.2f}s: {context}")
return context, token_count, generation_time
except Exception as e:
logger.error(f"Error generating chunk context: {e}")
# Return a basic context on error
basic_context = f"[Conversation chunk {chunk.chunk_index + 1}, Rounds {chunk.start_round}-{chunk.end_round}]"
return basic_context, 0, time.time() - start_time
def get_statistics(self) -> Dict[str, Any]:
"""Get chunking statistics"""
stats = self.stats.copy()
# Calculate averages
if stats["contextual_chunks"] > 0:
stats["avg_context_tokens"] = stats["total_context_tokens"] / stats["contextual_chunks"]
stats["avg_generation_time"] = stats["total_generation_time"] / stats["contextual_chunks"]
else:
stats["avg_context_tokens"] = 0
stats["avg_generation_time"] = 0
# Cache efficiency
total_cache_ops = stats["cache_hits"] + stats["cache_misses"]
if total_cache_ops > 0:
stats["cache_hit_rate"] = stats["cache_hits"] / total_cache_ops
else:
stats["cache_hit_rate"] = 0
# Cost estimation (rough)
if self.llm_config.provider == "openai":
cost_per_1k = 0.03
else:
cost_per_1k = 0.01
stats["estimated_cost"] = (stats["total_context_tokens"] / 1000) * cost_per_1k
return stats
contextual_compare.py¶
"""离线对比：上下文化记忆块 vs 原始记忆块对『用户事实召回』的影响（实验 3-12）。
本模块是一个**完全离线、无需 API/外部服务**的对照实验，用于量化本章核心论点：
在把对话记忆块送入索引/嵌入之前，先为每个块生成一段『上下文前缀』（情境锚定），
能显著提升脱离上下文的孤立片段（如『好的，就订这个吧』）被正确召回的概率。
方法说明（诚实边界）：
- 生产管线用 LLM 逐块生成 context、并用神经嵌入 + 检索服务做稠密/混合检索（需 API Key）。
- 这里用**确定性的 BM25 词法检索**作为无需 API 的代理：对同一份 context，
分别度量『不拼接（plain）』与『拼接后再索引（contextual）』两种方式的召回。
变量只有『索引文本是否含 context 前缀』，因此结果直接反映上下文化本身的贡献。
- 指标：Recall@1、Recall@3、MRR（口径同书中 recall@k：前 k 个结果命中即算召回）。
数据集见同目录 memory_qa_eval.json（受控教学集，可自行替换 --dataset）。
"""
import argparse
import json
import math
import re
from pathlib import Path
from typing import Dict, List, Tuple
DEFAULT_DATASET = str(Path(__file__).parent / "memory_qa_eval.json")
_CJK = r"一-鿿"
def tokenize(text: str) -> List[str]:
"""CJK 感知的轻量分词：英文/数字按词，中文按『单字 + 相邻双字』。
对相邻双字（bigram）保留位置邻接关系，避免跨片段产生虚假二元组。
"""
text = text.lower()
tokens: List[str] = []
tokens.extend(re.findall(r"[a-z0-9]+", text))
for run in re.findall(f"[{_CJK}]+", text):
chars = list(run)
tokens.extend(chars)
for i in range(len(chars) - 1):
tokens.append(chars[i] + chars[i + 1])
return tokens
class BM25:
"""标准 BM25，纯 Python 实现，无第三方依赖。"""
def __init__(self, corpus_tokens: List[List[str]], k1: float = 1.5, b: float = 0.75):
self.k1 = k1
self.b = b
self.corpus = corpus_tokens
self.N = len(corpus_tokens)
self.doc_len = [len(d) for d in corpus_tokens]
self.avgdl = sum(self.doc_len) / self.N if self.N else 0.0
self.tf: List[Dict[str, int]] = []
df: Dict[str, int] = {}
for doc in corpus_tokens:
counts: Dict[str, int] = {}
for t in doc:
counts[t] = counts.get(t, 0) + 1
self.tf.append(counts)
for t in counts:
df[t] = df.get(t, 0) + 1
self.idf = {
t: math.log(1 + (self.N - n + 0.5) / (n + 0.5)) for t, n in df.items()
}
def score(self, query_tokens: List[str], idx: int) -> float:
counts = self.tf[idx]
dl = self.doc_len[idx]
s = 0.0
for t in query_tokens:
if t not in counts:
continue
idf = self.idf.get(t, 0.0)
freq = counts[t]
denom = freq + self.k1 * (1 - self.b + self.b * dl / self.avgdl)
s += idf * (freq * (self.k1 + 1)) / denom
return s
def rank(self, query_tokens: List[str]) -> List[int]:
scored = [(self.score(query_tokens, i), i) for i in range(self.N)]
# 稳定排序：分数降序，平局按原始下标升序
scored.sort(key=lambda x: (-x[0], x[1]))
return [i for _, i in scored]
def _gold_rank(ranked_ids: List[str], gold_id: str) -> int:
"""返回 gold 在排序结果中的名次（1-based）；未找到返回 0。"""
for pos, cid in enumerate(ranked_ids, start=1):
if cid == gold_id:
return pos
return 0
def evaluate(chunks: List[dict], queries: List[dict], mode: str) -> Tuple[dict, List[dict]]:
"""按指定模式索引并检索，返回聚合指标与逐条明细。
mode='plain' 索引文本 = chunk['text']
mode='contextual' 索引文本 = chunk['context'] + '\n' + chunk['text']
"""
ids = [c["id"] for c in chunks]
if mode == "plain":
docs = [c["text"] for c in chunks]
elif mode == "contextual":
docs = [f"{c.get('context','')}\n{c['text']}" for c in chunks]
else:
raise ValueError(f"unknown mode: {mode}")
bm25 = BM25([tokenize(d) for d in docs])
r1 = r3 = 0
mrr = 0.0
details = []
for q in queries:
ranked = [ids[i] for i in bm25.rank(tokenize(q["q"]))]
rank = _gold_rank(ranked, q["gold"])
hit1 = 1 if rank == 1 else 0
hit3 = 1 if 1 <= rank <= 3 else 0
r1 += hit1
r3 += hit3
mrr += (1.0 / rank) if rank else 0.0
details.append({"q": q["q"], "gold": q["gold"], "rank": rank,
"hit@1": hit1, "hit@3": hit3, "top": ranked[:3]})
n = len(queries)
metrics = {
"recall@1": r1 / n,
"recall@3": r3 / n,
"mrr": mrr / n,
"n_queries": n,
}
return metrics, details
def run_comparison(dataset_path: str, output_path: str = None, verbose: bool = True) -> dict:
data = json.loads(Path(dataset_path).read_text(encoding="utf-8"))
chunks, queries = data["chunks"], data["queries"]
plain_m, plain_d = evaluate(chunks, queries, "plain")
ctx_m, ctx_d = evaluate(chunks, queries, "contextual")
if verbose:
print("=" * 68)
print("实验 3-12 · 上下文化记忆块对用户事实召回的影响（离线 BM25 代理）")
print(f"数据集: {dataset_path}")
print(f"记忆块: {len(chunks)} 查询: {len(queries)}")
print("=" * 68)
print(f"{'方法':<28}{'Recall@1':>10}{'Recall@3':>10}{'MRR':>10}")
print("-" * 68)
print(f"{'Plain（直接索引原始块）':<24}{plain_m['recall@1']:>10.3f}"
f"{plain_m['recall@3']:>10.3f}{plain_m['mrr']:>10.3f}")
print(f"{'Contextual（上下文化后索引）':<22}{ctx_m['recall@1']:>10.3f}"
f"{ctx_m['recall@3']:>10.3f}{ctx_m['mrr']:>10.3f}")
print("-" * 68)
d1 = ctx_m["recall@1"] - plain_m["recall@1"]
d3 = ctx_m["recall@3"] - plain_m["recall@3"]
dm = ctx_m["mrr"] - plain_m["mrr"]
print(f"{'提升（Δ）':<26}{d1:>+10.3f}{d3:>+10.3f}{dm:>+10.3f}")
print("=" * 68)
print("\n逐查询名次（gold 在检索结果中的位次，越小越好；0 表示未召回）：")
print(f"{'查询':<38}{'Plain':>8}{'Ctx':>8}")
print("-" * 68)
pd = {x["q"]: x["rank"] for x in plain_d}
for x in ctx_d:
q = x["q"][:36]
print(f"{q:<38}{pd[x['q']]:>8}{x['rank']:>8}")
print("=" * 68)
print("说明：孤立片段（如『好的，就订这个吧』）在 Plain 下缺乏可检索信号，")
print("上下文化后被『锚定』回其情境，因而召回名次明显提升。")
result = {
"dataset": dataset_path,
"n_chunks": len(chunks),
"plain": plain_m,
"contextual": ctx_m,
"delta": {
"recall@1": ctx_m["recall@1"] - plain_m["recall@1"],
"recall@3": ctx_m["recall@3"] - plain_m["recall@3"],
"mrr": ctx_m["mrr"] - plain_m["mrr"],
},
"plain_details": plain_d,
"contextual_details": ctx_d,
}
if output_path:
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
Path(output_path).write_text(json.dumps(result, ensure_ascii=False, indent=2),
encoding="utf-8")
if verbose:
print(f"\n结果已保存至: {output_path}")
return result
def single_query(dataset_path: str, query: str, top_k: int = 3, verbose: bool = True) -> dict:
"""针对单条查询，离线对比 plain 与 contextual 两种索引下的 Top-K 检索结果。"""
data = json.loads(Path(dataset_path).read_text(encoding="utf-8"))
chunks = data["chunks"]
ids = [c["id"] for c in chunks]
qt = tokenize(query)
out = {"query": query}
for mode in ("plain", "contextual"):
if mode == "plain":
docs = [c["text"] for c in chunks]
else:
docs = [f"{c.get('context','')}\n{c['text']}" for c in chunks]
bm25 = BM25([tokenize(d) for d in docs])
ranked = bm25.rank(qt)
out[mode] = [{"id": ids[i], "score": round(bm25.score(qt, i), 4)}
for i in ranked[:top_k]]
if verbose:
print("=" * 60)
print(f"查询: {query}")
print("=" * 60)
for mode in ("plain", "contextual"):
label = "Plain（原始块）" if mode == "plain" else "Contextual（上下文化）"
print(f"\n[{label}] Top-{top_k}:")
for rank, item in enumerate(out[mode], 1):
print(f" {rank}. {item['id']:<18} score={item['score']}")
print("=" * 60)
return out
def build_arg_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="离线对比上下文化记忆块 vs 原始记忆块的用户事实召回（实验 3-12，无需 API）",
)
p.add_argument("--dataset", default=DEFAULT_DATASET,
help="记忆问答对照集 JSON 路径（默认：memory_qa_eval.json）")
p.add_argument("--output", default=None,
help="将对比结果（含逐查询明细）保存为 JSON 的路径（默认不保存）")
p.add_argument("--quiet", action="store_true", help="仅输出结果 JSON，不打印表格")
return p
def main():
args = build_arg_parser().parse_args()
run_comparison(args.dataset, output_path=args.output, verbose=not args.quiet)
if __name__ == "__main__":
main()
contextual_evaluator.py¶
"""Contextual Memory Evaluator with Dual Memory System
This evaluator tests the combined system of:
1. Contextual RAG for conversation chunks
2. Advanced JSON cards for structured facts
"""
import os
import json
import logging
import time
import uuid
import sys
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from config import Config
from contextual_indexer import ContextualMemoryIndexer
from contextual_agent import ContextualUserMemoryAgent
from advanced_memory_manager import AdvancedMemoryManager, AdvancedMemoryCard
from chunker import ConversationChunker, ConversationChunk
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Import LLM evaluator from user-memory-evaluation project
LLMEvaluator = None
EvalTestCase = None
ConversationHistory = None
EvalMessage = None
MessageRole = None
# Try to import from week2/user-memory-evaluation
eval_project_path = Path(__file__).parent.parent.parent / "week2" / "user-memory-evaluation"
if eval_project_path.exists():
sys.path.insert(0, str(eval_project_path))
try:
from llm_evaluator import LLMEvaluator
from test_case import TestCase as EvalTestCase, ConversationHistory, Message as EvalMessage, MessageRole
logger.info("Successfully imported LLM evaluation modules")
except ImportError as e:
logger.warning(f"Could not import LLM evaluation modules: {e}")
logger.info("LLM evaluation will not be available")
else:
logger.warning(f"user-memory-evaluation project not found at {eval_project_path}")
logger.info("LLM evaluation will not be available")
@dataclass
class TestCase:
"""Enhanced test case with support for advanced memory"""
test_id: str
category: str # layer1, layer2, layer3
title: str
description: str
conversation_histories: List[Dict[str, Any]] # Conversation data directly from YAML
user_question: str
evaluation_criteria: str
expected_behavior: Optional[str] = None
expected_memory_cards: List[Dict[str, Any]] = field(default_factory=list)
metadata: Dict[str, Any] = field(default_factory=dict)
@classmethod
def from_yaml(cls, data: Dict[str, Any], test_dir: Path) -> 'TestCase':
"""Create test case from YAML data"""
# The conversation_histories field directly contains the conversation data
conversation_histories = data.get('conversation_histories', [])
return cls(
test_id=data.get('test_id', data.get('id', '')),
category=data.get('category', ''),
title=data.get('title', ''),
description=data.get('description', ''),
conversation_histories=conversation_histories,
user_question=data.get('user_question', ''),
evaluation_criteria=data.get('evaluation_criteria', ''),
expected_behavior=data.get('expected_behavior'),
expected_memory_cards=data.get('expected_memory_cards', []),
metadata=data.get('metadata', {})
)
@dataclass
class EvaluationResult:
"""Enhanced evaluation result with dual memory tracking"""
test_id: str
success: bool
agent_answer: str
evaluation_criteria: str
iterations: int
tool_calls: int
memory_cards_used: List[str]
chunks_retrieved: List[str]
contextual_chunks_count: int
processing_time: float
indexing_time: float
context_generation_time: float
llm_evaluation: Optional[Dict[str, Any]] = None
error: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {
"test_id": self.test_id,
"success": self.success,
"agent_answer": self.agent_answer,
"evaluation_criteria": self.evaluation_criteria,
"iterations": self.iterations,
"tool_calls": self.tool_calls,
"memory_cards_used": self.memory_cards_used,
"chunks_retrieved": self.chunks_retrieved,
"contextual_chunks_count": self.contextual_chunks_count,
"processing_time": self.processing_time,
"indexing_time": self.indexing_time,
"context_generation_time": self.context_generation_time,
"llm_evaluation": self.llm_evaluation,
"error": self.error
}
class ContextualMemoryEvaluator:
"""Evaluates the dual memory system on test cases"""
def __init__(self, config: Optional[Config] = None):
"""
Initialize the evaluator
Args:
config: Configuration object
"""
self.config = config or Config.from_env()
self.test_cases: Dict[str, TestCase] = {}
self.results: Dict[str, EvaluationResult] = {}
# Initialize components
self.chunker = ConversationChunker(self.config.chunking)
self.indexer: Optional[ContextualMemoryIndexer] = None
self.agent: Optional[ContextualUserMemoryAgent] = None
# Initialize LLM evaluator if available
self.llm_evaluator = None
logger.info("Checking LLM Evaluator availability...")
logger.info(f"LLMEvaluator module: {LLMEvaluator}")
logger.info(f"EvalTestCase module: {EvalTestCase}")
if LLMEvaluator:
try:
logger.info("Attempting to initialize LLM Evaluator...")
self.llm_evaluator = LLMEvaluator()
logger.info("✅ LLM Evaluator initialized successfully for automatic evaluation")
except Exception as e:
logger.warning(f"Could not initialize LLM evaluator: {e}")
logger.info("Automatic LLM evaluation will be skipped")
else:
logger.info("LLM Evaluator not available - automatic evaluation will be skipped")
# Evaluation framework path
self.eval_framework_path = Path("../user-memory-evaluation/test_cases")
if not self.eval_framework_path.exists():
self.eval_framework_path = Path("../../week2/user-memory-evaluation/test_cases")
logger.info("Initialized ContextualMemoryEvaluator")
def load_test_cases(self, category: Optional[str] = None) -> List[str]:
"""
Load test cases from the evaluation framework
Args:
category: Optional category filter (layer1, layer2, layer3)
Returns:
List of loaded test case IDs
"""
import yaml
test_cases = []
test_dirs = ["layer1", "layer2", "layer3"] if not category else [category]
for test_dir in test_dirs:
dir_path = self.eval_framework_path / test_dir
if not dir_path.exists():
logger.warning(f"Test directory not found: {dir_path}")
continue
# Load all YAML files
for yaml_file in dir_path.glob("*.yaml"):
try:
with open(yaml_file, 'r', encoding='utf-8') as f:
data = yaml.safe_load(f)
test_case = TestCase.from_yaml(data, self.eval_framework_path)
test_case.category = test_dir # Ensure category is set
self.test_cases[test_case.test_id] = test_case
test_cases.append(test_case.test_id)
except Exception as e:
logger.error(f"Error loading test case {yaml_file}: {e}")
logger.info(f"Loaded {len(test_cases)} test cases")
return test_cases
def _prepare_memory_cards(self, test_case: TestCase) -> List[AdvancedMemoryCard]:
"""
Generate memory cards for a test case
For Layer 3 tests, this creates cards that enable proactive service
"""
cards = []
# Generate cards based on test metadata
if test_case.category == "layer3":
# Layer 3 needs structured facts for proactive service
if "travel" in test_case.test_id.lower():
cards.append(AdvancedMemoryCard(
category="travel",
card_key="tokyo_trip",
backstory="User booked a trip to Tokyo in previous conversations",
date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"destination": "Tokyo, Japan",
"departure_date": "2025-01-25",
"return_date": "2025-02-01",
"purpose": "business conference"
}
))
cards.append(AdvancedMemoryCard(
category="travel",
card_key="passport_jessica",
backstory="User's passport expiration was mentioned when booking travel",
date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"passport_number": "XXXXX1234",
"expiration_date": "2025-02-18",
"issuing_country": "USA",
"needs_renewal": True
}
))
# Add any expected cards from test case
for card_data in test_case.expected_memory_cards:
cards.append(AdvancedMemoryCard(
category=card_data.get("category", "general"),
card_key=card_data.get("key", str(uuid.uuid4())),
backstory=card_data.get("backstory", "From test case"),
date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
person=card_data.get("person", "User"),
relationship=card_data.get("relationship", "primary"),
data=card_data.get("data", {})
))
return cards
def evaluate_test_case(self, test_id: str) -> EvaluationResult:
"""
Evaluate a single test case with dual memory system
Args:
test_id: Test case ID
Returns:
Evaluation result
"""
if test_id not in self.test_cases:
raise ValueError(f"Test case {test_id} not found")
test_case = self.test_cases[test_id]
logger.info(f"Evaluating test case: {test_id} - {test_case.title}")
start_time = time.time()
indexing_start = time.time()
try:
# Step 1: Initialize indexer with user ID
user_id = f"test_user_{test_id}"
self.indexer = ContextualMemoryIndexer(
user_id=user_id,
index_config=self.config.index,
chunking_config=self.config.chunking,
use_contextual=True # Enable contextual chunking
)
# Step 2: Process conversations from test case
all_chunks = []
for conv_data in test_case.conversation_histories:
# Extract conversation data from the test case
conv_id = conv_data.get('conversation_id', f'conv_{test_id}')
messages = conv_data.get('messages', [])
# Chunk the conversation
chunks = self.chunker.chunk_conversation(
messages=messages,
conversation_id=conv_id,
test_id=test_id
)
all_chunks.extend(chunks)
logger.info(f"Created {len(all_chunks)} basic chunks")
# Step 3: Process with contextual chunking and indexing
processing_result = self.indexer.process_conversation_history(
chunks=all_chunks,
conversation_id=test_id,
generate_summary_cards=True # Generate cards from conversations
)
# Step 4: Add pre-defined memory cards for the test
memory_cards = self._prepare_memory_cards(test_case)
for card in memory_cards:
self.indexer.memory_manager.add_card(card)
# Debug: Print all memory cards
logger.info("="*60)
logger.info("DEBUG: All Memory Cards in System")
logger.info("="*60)
# Access cards directly from categories attribute
for category, cards in self.indexer.memory_manager.categories.items():
for card_key, card in cards.items():
logger.info(f"\n[{category}.{card_key}]")
logger.info(json.dumps(card.to_dict(), indent=2, ensure_ascii=False))
total_cards = sum(len(cards) for cards in self.indexer.memory_manager.categories.values())
logger.info(f"\nTotal Memory Cards: {total_cards}")
logger.info("="*60)
indexing_time = time.time() - indexing_start
# Step 5: Initialize agent with dual memory
self.agent = ContextualUserMemoryAgent(
indexer=self.indexer,
memory_manager=self.indexer.memory_manager,
config=self.config
)
# Step 6: Answer the question
trajectory = self.agent.answer_question(
question=test_case.user_question,
test_id=test_id,
max_iterations=self.config.evaluation.max_iterations,
stream=False
)
processing_time = time.time() - start_time
# Get context generation statistics
chunker_stats = self.indexer.contextual_chunker.get_statistics()
context_gen_time = chunker_stats.get("total_generation_time", 0)
# Step 7: Evaluate the answer with LLM (automatic if available)
llm_evaluation = None
logger.info("="*60)
logger.info("LLM Judge Evaluation")
logger.info("="*60)
if not self.llm_evaluator:
logger.warning("LLM Judge not available - skipping automatic evaluation")
logger.info("To enable LLM Judge, ensure the llm_evaluator module is properly imported")
elif not trajectory.final_answer:
logger.warning("No final answer from agent - cannot evaluate")
else:
logger.info("Running LLM Judge evaluation...")
if self.llm_evaluator and trajectory.final_answer:
try:
# Convert test case to evaluation format
eval_histories = []
for hist in test_case.conversation_histories:
eval_messages = []
for msg in hist.get("messages", []):
role_str = msg.get("role", "user")
role = MessageRole.USER if role_str == "user" else MessageRole.ASSISTANT
eval_messages.append(EvalMessage(
role=role,
content=msg.get("content", "")
))
eval_histories.append(ConversationHistory(
conversation_id=hist.get("conversation_id", f"conv_{uuid.uuid4().hex[:8]}"),
messages=eval_messages,
metadata=hist.get("metadata", {})
))
# Create evaluation test case
eval_test_case = EvalTestCase(
test_id=test_case.test_id,
category=test_case.category,
title=test_case.title,
description=test_case.description,
conversation_histories=eval_histories,
user_question=test_case.user_question,
evaluation_criteria=test_case.evaluation_criteria if test_case.evaluation_criteria else "The agent should provide a relevant and accurate response based on the conversation history.",
expected_behavior=test_case.expected_behavior
)
# Run LLM evaluation
llm_result = self.llm_evaluator.evaluate(
test_case=eval_test_case,
agent_response=trajectory.final_answer,
extracted_memory=None
)
llm_evaluation = {
"reward": llm_result.reward,
"passed": llm_result.passed if llm_result.passed is not None else llm_result.reward >= 0.6,
"reasoning": llm_result.reasoning,
"required_info_found": llm_result.required_info_found if hasattr(llm_result, 'required_info_found') else {},
"suggestions": llm_result.suggestions if hasattr(llm_result, 'suggestions') else None
}
# Log evaluation results
logger.info("-"*60)
logger.info(f"LLM Evaluation Reward: {llm_result.reward:.3f}/1.000")
logger.info(f"Passed: {'Yes' if llm_evaluation['passed'] else 'No'}")
logger.info("-"*60)
logger.info(f"Evaluation Reasoning:")
logger.info(llm_result.reasoning)
logger.info("-"*60)
if llm_result.required_info_found:
logger.info("Required Information Found:")
for key, found in llm_result.required_info_found.items():
status = "✓" if found else "✗"
logger.info(f" {status} {key}")
logger.info("-"*60)
except Exception as e:
logger.warning(f"LLM evaluation failed: {e}")
logger.debug("Full error:", exc_info=True)
else:
# Fallback: Use direct LLM API for evaluation if module not available
if trajectory.final_answer:
logger.info("Attempting fallback LLM evaluation...")
llm_evaluation = self._fallback_llm_evaluation(test_case, trajectory.final_answer)
# Create result
result = EvaluationResult(
test_id=test_id,
success=trajectory.success,
agent_answer=trajectory.final_answer or "",
evaluation_criteria=test_case.evaluation_criteria,
iterations=len(trajectory.iterations),
tool_calls=len(trajectory.tool_calls),
memory_cards_used=trajectory.memory_cards_used,
chunks_retrieved=trajectory.chunks_retrieved,
contextual_chunks_count=processing_result.get("contextual_chunks", 0),
processing_time=processing_time,
indexing_time=indexing_time,
context_generation_time=context_gen_time,
llm_evaluation=llm_evaluation
)
# Check if answer meets criteria
# Use LLM evaluation result if available, otherwise fall back to keyword check
if llm_evaluation and 'passed' in llm_evaluation:
result.success = llm_evaluation['passed']
logger.info(f"Using LLM evaluation result: {'Success' if result.success else 'Failed'}")
elif self._check_answer_criteria(trajectory.final_answer, test_case.evaluation_criteria):
result.success = True
logger.info(f"Using keyword-based evaluation: Success")
else:
logger.info(f"Using keyword-based evaluation: Failed")
self.results[test_id] = result
logger.info(f"Evaluation complete for {test_id}: {'Success' if result.success else 'Failed'}")
return result
except Exception as e:
logger.error(f"Error evaluating test case {test_id}: {e}")
result = EvaluationResult(
test_id=test_id,
success=False,
agent_answer="",
evaluation_criteria=test_case.evaluation_criteria,
iterations=0,
tool_calls=0,
memory_cards_used=[],
chunks_retrieved=[],
contextual_chunks_count=0,
processing_time=time.time() - start_time,
indexing_time=indexing_time if 'indexing_time' in locals() else 0,
context_generation_time=0,
error=str(e)
)
self.results[test_id] = result
return result
def _load_conversations(self, conv_file: str) -> Dict[str, List[Dict[str, str]]]:
"""Load conversations from a JSON file"""
conversations = {}
try:
with open(conv_file, 'r', encoding='utf-8') as f:
data = json.load(f)
# Handle different formats
if isinstance(data, dict):
if "conversations" in data:
conversations = data["conversations"]
else:
# Assume it's already in the right format
conversations = data
elif isinstance(data, list):
# Convert list to dict with generated IDs
for i, conv in enumerate(data):
conversations[f"conv_{i}"] = conv
return conversations
except Exception as e:
logger.error(f"Error loading conversations from {conv_file}: {e}")
return {}
def _check_answer_criteria(self, answer: Optional[str], criteria: str) -> bool:
"""Simple check if answer meets criteria"""
if not answer:
return False
answer_lower = answer.lower()
criteria_lower = criteria.lower()
# Extract key terms from criteria
key_terms = []
for word in criteria_lower.split():
if len(word) > 4 and word not in ["should", "must", "need", "have"]:
key_terms.append(word)
# Check if key terms appear in answer
matches = sum(1 for term in key_terms if term in answer_lower)
return matches >= min(3, len(key_terms) // 2)
def _evaluate_with_llm(self, test_case: TestCase, agent_answer: Optional[str]) -> Dict[str, Any]:
"""Use LLM to evaluate if the answer meets criteria"""
if not agent_answer:
return {
"passed": False,
"reasoning": "No answer provided",
"reward": 0.0
}
# This is a simplified version - implement full LLM evaluation as needed
return {
"passed": self._check_answer_criteria(agent_answer, test_case.evaluation_criteria),
"reasoning": "Basic criteria check",
"reward": 0.5
}
def _fallback_llm_evaluation(self, test_case: TestCase, agent_answer: str) -> Dict[str, Any]:
"""Fallback LLM evaluation using direct API call when llm_evaluator module isn't available"""
try:
from openai import OpenAI
config = Config.from_env()
client_config, model = config.llm.get_client_config()
base_url = client_config.pop("base_url", None)
if base_url:
client = OpenAI(base_url=base_url, **client_config)
else:
client = OpenAI(**client_config)
# Create evaluation prompt
eval_prompt = f"""Evaluate the agent's response based on the test criteria.
Test Question: {test_case.user_question}
Agent's Answer: {agent_answer}
Evaluation Criteria: {test_case.evaluation_criteria}
Provide a JSON evaluation with:
1. "reward": A score from 0.0 to 1.0 (0.6+ is passing)
2. "passed": Boolean indicating if the answer meets the criteria
3. "reasoning": Explanation for the score (2-3 sentences)
4. "required_info_found": Object with boolean values for each required piece of information
Respond with valid JSON only."""
response = client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": "You are an evaluation judge. Evaluate if the agent's answer correctly addresses the user's question based on the criteria."},
{"role": "user", "content": eval_prompt}
],
temperature=_reasoning_safe_temperature(model, 0.1),
response_format={"type": "json_object"}
)
result = json.loads(response.choices[0].message.content)
# Ensure required fields
if 'reward' not in result:
result['reward'] = 0.5
if 'passed' not in result:
result['passed'] = result['reward'] >= 0.6
if 'reasoning' not in result:
result['reasoning'] = "Evaluation completed"
if 'required_info_found' not in result:
result['required_info_found'] = {}
# Log evaluation results
logger.info("="*60)
logger.info("Fallback LLM Evaluation Results")
logger.info("="*60)
logger.info(f"Reward: {result['reward']:.3f}/1.000")
logger.info(f"Passed: {'Yes' if result['passed'] else 'No'}")
logger.info(f"Reasoning: {result['reasoning']}")
if result['required_info_found']:
logger.info("Required Information Found:")
for key, found in result['required_info_found'].items():
status = "✓" if found else "✗"
logger.info(f" {status} {key}")
logger.info("="*60)
return result
except Exception as e:
logger.error(f"Fallback LLM evaluation failed: {e}")
return None
def evaluate_batch(self, test_ids: Optional[List[str]] = None) -> Dict[str, EvaluationResult]:
"""Evaluate multiple test cases"""
if test_ids is None:
test_ids = list(self.test_cases.keys())
for test_id in test_ids:
try:
self.evaluate_test_case(test_id)
except Exception as e:
logger.error(f"Failed to evaluate {test_id}: {e}")
return self.results
def generate_report(self) -> str:
"""Generate evaluation report"""
if not self.results:
return "No evaluation results available"
report = ["=" * 80]
report.append("CONTEXTUAL MEMORY EVALUATION REPORT")
report.append("=" * 80)
report.append(f"Generated: {datetime.now().isoformat()}")
report.append(f"Total Test Cases: {len(self.results)}")
report.append("")
# Summary statistics
successful = sum(1 for r in self.results.values() if r.success)
report.append(f"Success Rate: {successful}/{len(self.results)} ({100*successful/len(self.results):.1f}%)")
report.append("")
# Statistics by category
categories = {}
for test_id, result in self.results.items():
if test_id in self.test_cases:
cat = self.test_cases[test_id].category
if cat not in categories:
categories[cat] = {"total": 0, "success": 0}
categories[cat]["total"] += 1
if result.success:
categories[cat]["success"] += 1
report.append("Results by Category:")
for cat in sorted(categories.keys()):
stats = categories[cat]
rate = 100 * stats["success"] / stats["total"] if stats["total"] > 0 else 0
report.append(f" {cat}: {stats['success']}/{stats['total']} ({rate:.1f}%)")
report.append("")
# Average metrics
avg_iterations = sum(r.iterations for r in self.results.values()) / len(self.results)
avg_tool_calls = sum(r.tool_calls for r in self.results.values()) / len(self.results)
avg_time = sum(r.processing_time for r in self.results.values()) / len(self.results)
avg_context_time = sum(r.context_generation_time for r in self.results.values()) / len(self.results)
report.append("Average Metrics:")
report.append(f" Iterations: {avg_iterations:.1f}")
report.append(f" Tool Calls: {avg_tool_calls:.1f}")
report.append(f" Processing Time: {avg_time:.2f}s")
report.append(f" Context Generation Time: {avg_context_time:.2f}s")
report.append("")
# Memory usage statistics
total_cards_used = sum(len(r.memory_cards_used) for r in self.results.values())
total_chunks_retrieved = sum(len(r.chunks_retrieved) for r in self.results.values())
report.append("Memory Usage:")
report.append(f" Total Memory Cards Used: {total_cards_used}")
report.append(f" Total Chunks Retrieved: {total_chunks_retrieved}")
report.append(f" Avg Cards per Query: {total_cards_used/len(self.results):.1f}")
report.append(f" Avg Chunks per Query: {total_chunks_retrieved/len(self.results):.1f}")
report.append("")
# Individual test results
report.append("-" * 80)
report.append("INDIVIDUAL TEST RESULTS")
report.append("-" * 80)
for test_id, result in sorted(self.results.items()):
test_case = self.test_cases.get(test_id)
report.append(f"\nTest: {test_id}")
if test_case:
report.append(f" Title: {test_case.title}")
report.append(f" Category: {test_case.category}")
report.append(f" Status: {'✓ Success' if result.success else '✗ Failed'}")
report.append(f" Iterations: {result.iterations}")
report.append(f" Tool Calls: {result.tool_calls}")
report.append(f" Memory Cards Used: {len(result.memory_cards_used)}")
report.append(f" Chunks Retrieved: {len(result.chunks_retrieved)}")
report.append(f" Contextual Chunks: {result.contextual_chunks_count}")
report.append(f" Processing Time: {result.processing_time:.2f}s")
if result.error:
report.append(f" Error: {result.error}")
return "\n".join(report)
def save_results(self, output_file: str):
"""Save evaluation results to file"""
results_data = {
"timestamp": datetime.now().isoformat(),
"config": {
"use_contextual": True,
"llm_provider": self.config.llm.provider,
"llm_model": self.config.llm.model,
"chunking_strategy": self.config.chunking.strategy.value,
"index_mode": self.config.index.mode.value
},
"summary": {
"total_tests": len(self.results),
"successful": sum(1 for r in self.results.values() if r.success),
"failed": sum(1 for r in self.results.values() if not r.success)
},
"results": {
test_id: result.to_dict()
for test_id, result in self.results.items()
}
}
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(results_data, f, indent=2, ensure_ascii=False)
logger.info(f"Results saved to {output_file}")
contextual_indexer.py¶
"""Contextual RAG Indexer with Advanced Memory Cards
This module combines contextual chunking for conversation histories with
advanced JSON cards for structured user memory.
"""
import os
import json
import logging
import requests
import time
from datetime import datetime
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path
from config import IndexConfig, IndexMode, ChunkingConfig
from chunker import ConversationChunk
from contextual_chunking import ContextualConversationChunk, ContextualConversationChunker
from advanced_memory_manager import AdvancedMemoryManager, AdvancedMemoryCard
from indexer import SearchResult
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class ContextualSearchResult(SearchResult):
"""Enhanced search result with contextual information"""
context: str = ""
contextual_chunk: Optional[ContextualConversationChunk] = None
def to_dict(self) -> Dict[str, Any]:
result = super().to_dict()
result["context"] = self.context
if self.contextual_chunk:
result["contextual_info"] = {
"context": self.contextual_chunk.context,
"context_tokens": self.contextual_chunk.context_tokens,
"generation_time": self.contextual_chunk.generation_time
}
return result
class ContextualMemoryIndexer:
"""
Dual-layer memory indexer combining:
1. Contextual RAG for conversation chunks
2. Advanced JSON cards for structured facts
"""
def __init__(self,
user_id: str,
index_config: Optional[IndexConfig] = None,
chunking_config: Optional[ChunkingConfig] = None,
use_contextual: bool = True):
"""
Initialize the contextual memory indexer.
Args:
user_id: User identifier
index_config: Index configuration
chunking_config: Chunking configuration
use_contextual: Whether to use contextual chunking
"""
self.user_id = user_id
self.index_config = index_config or IndexConfig()
self.chunking_config = chunking_config or ChunkingConfig()
self.use_contextual = use_contextual
# Initialize contextual chunker
self.contextual_chunker = ContextualConversationChunker(
chunking_config=self.chunking_config,
use_contextual=use_contextual
)
# Initialize advanced memory manager
self.memory_manager = AdvancedMemoryManager(
user_id=user_id,
storage_dir=self.index_config.chunk_store_path
)
# Storage for chunks
self.contextual_chunks: Dict[str, ContextualConversationChunk] = {}
self.doc_id_mapping: Dict[str, str] = {}
# Retrieval pipeline URL
self.retrieval_url = "http://localhost:4242"
# Create directories
Path(self.index_config.index_path).parent.mkdir(parents=True, exist_ok=True)
Path(self.index_config.chunk_store_path).mkdir(parents=True, exist_ok=True)
# Statistics
self.stats = {
"chunks_indexed": 0,
"memory_cards": 0,
"context_generation_time": 0.0,
"indexing_time": 0.0
}
self._check_retrieval_pipeline()
logger.info(f"Initialized ContextualMemoryIndexer for user {user_id}")
def _check_retrieval_pipeline(self):
"""Check if the retrieval pipeline service is available"""
try:
response = requests.get(f"{self.retrieval_url}/health", timeout=2)
if response.status_code == 200:
logger.info("✓ Retrieval pipeline service is available")
else:
logger.warning(f"Retrieval pipeline returned status {response.status_code}")
except requests.exceptions.RequestException as e:
logger.warning(f"Retrieval pipeline service not available at {self.retrieval_url}: {e}")
logger.info("Note: The retrieval pipeline needs to be running. Start it with:")
logger.info(" cd projects/week3/retrieval-pipeline && python api_server.py")
def process_conversation_history(self,
chunks: List[ConversationChunk],
conversation_id: str,
generate_summary_cards: bool = True) -> Dict[str, Any]:
"""
Process conversation history with both contextual chunking and memory cards.
Args:
chunks: Basic conversation chunks
conversation_id: Conversation identifier
generate_summary_cards: Whether to generate summary cards
Returns:
Processing results
"""
start_time = time.time()
results = {
"conversation_id": conversation_id,
"contextual_chunks": 0,
"memory_cards_before": len(self.memory_manager.categories),
"memory_cards_after": 0,
"processing_time": 0
}
# Step 1: Generate contextual chunks
logger.info(f"Processing {len(chunks)} chunks for conversation {conversation_id}")
if self.use_contextual:
contextual_chunks = self.contextual_chunker.contextualize_chunks(chunks)
logger.info(f"Generated {len(contextual_chunks)} contextual chunks")
else:
# Convert to contextual chunks without context
contextual_chunks = [
ContextualConversationChunk.from_basic_chunk(chunk)
for chunk in chunks
]
# Store contextual chunks
for chunk in contextual_chunks:
self.contextual_chunks[chunk.chunk_id] = chunk
results["contextual_chunks"] = len(contextual_chunks)
# Step 2: Index contextual chunks
self._index_contextual_chunks(contextual_chunks)
# Step 3: Generate summary cards if requested
if generate_summary_cards:
summary_cards = self._generate_summary_cards(chunks, conversation_id)
for card in summary_cards:
self.memory_manager.add_card(card)
logger.info(f"Generated {len(summary_cards)} summary cards")
results["new_summary_cards"] = len(summary_cards)
# Step 4: Create conversation summary linking to memory cards
conversation_summary = self.memory_manager.summarize_for_conversation(conversation_id)
results["memory_cards_after"] = sum(len(cards) for cards in self.memory_manager.categories.values())
results["conversation_summary"] = conversation_summary
# Update statistics
results["processing_time"] = time.time() - start_time
self.stats["chunks_indexed"] += len(contextual_chunks)
self.stats["memory_cards"] = results["memory_cards_after"]
self.stats["indexing_time"] += results["processing_time"]
logger.info(f"Processing complete for conversation {conversation_id} in {results['processing_time']:.2f}s")
return results
def _index_contextual_chunks(self, chunks: List[ContextualConversationChunk]):
"""Index contextual chunks in the retrieval pipeline"""
documents = []
for chunk in chunks:
# Use contextualized text for indexing
doc = {
"text": chunk.contextualized_text,
"metadata": {
"doc_id": chunk.chunk_id,
"test_id": chunk.test_id,
"conversation_id": chunk.conversation_id,
"chunk_index": chunk.chunk_index,
"start_round": chunk.start_round,
"end_round": chunk.end_round,
"has_context": bool(chunk.context),
"context_preview": chunk.context[:100] if chunk.context else "",
**chunk.metadata
}
}
documents.append(doc)
# Send to retrieval pipeline
try:
indexed_count = 0
for doc in documents:
try:
response = requests.post(
f"{self.retrieval_url}/index",
json=doc
)
if response.status_code == 200:
result = response.json()
generated_doc_id = result.get("doc_id")
our_chunk_id = doc.get("metadata", {}).get("doc_id")
if generated_doc_id and our_chunk_id:
self.doc_id_mapping[generated_doc_id] = our_chunk_id
indexed_count += 1
except Exception as e:
logger.warning(f"Error indexing document: {e}")
logger.info(f"Indexed {indexed_count}/{len(documents)} contextual chunks")
except Exception as e:
logger.error(f"Error connecting to retrieval pipeline: {e}")
def _generate_summary_cards(self,
chunks: List[ConversationChunk],
conversation_id: str) -> List[AdvancedMemoryCard]:
"""
Generate summary cards from conversation chunks using LLM extraction.
Uses an LLM to analyze conversation chunks and extract structured memory cards
following the Advanced JSON Cards format from week2/user-memory.
"""
cards = []
if not chunks:
return cards
# Combine all conversation text
full_text = "\n".join([chunk.to_text() for chunk in chunks])
# Use LLM to extract structured memory cards
try:
# Initialize LLM client if needed
if not hasattr(self, '_llm_client'):
from openai import OpenAI
from config import Config
config = Config.from_env()
client_config, model = config.llm.get_client_config()
base_url = client_config.pop("base_url", None)
if base_url:
self._llm_client = OpenAI(base_url=base_url, **client_config)
else:
self._llm_client = OpenAI(**client_config)
self._llm_model = model
# Create extraction prompt
extraction_prompt = f"""Analyze this conversation and extract structured memory cards.
Conversation:
{full_text}
Extract any important information into memory cards. Each card MUST include:
- category: The type of information (e.g., 'personal', 'financial', 'medical', 'travel', 'family', 'work')
- card_key: A unique identifier for this card
- backstory: Context about when/why this information was learned (1-2 sentences)
- date_created: Current timestamp (YYYY-MM-DD HH:MM:SS)
- person: Who this relates to (e.g., "John Smith (primary)", "Sarah Smith (daughter)")
- relationship: Role/relationship (e.g., "primary account holder", "family member")
- Additional relevant fields based on the information type
Respond with a JSON array of memory cards. If no significant information found, return empty array [].
Example format:
[{{
"category": "financial",
"card_key": "bank_account_primary",
"backstory": "User shared their banking details while setting up automatic bill payments",
"date_created": "2024-01-15 10:30:00",
"person": "John Smith (primary)",
"relationship": "primary account holder",
"bank_name": "Chase Bank",
"account_type": "checking",
"account_ending": "4567"
}}]
Extract memory cards:"""
# Call LLM for extraction
response = self._llm_client.chat.completions.create(
model=self._llm_model,
messages=[
{"role": "system", "content": "You are a memory extraction assistant. Extract structured information from conversations into memory cards."},
{"role": "user", "content": extraction_prompt}
],
temperature=_reasoning_safe_temperature(self._llm_model, 0.3),
response_format={"type": "json_object"}
)
# Parse response
content = response.choices[0].message.content
extracted_data = json.loads(content) if content else {}
# Handle both single card and array of cards
if isinstance(extracted_data, dict):
# Check if it's a wrapper with 'cards' key
if 'cards' in extracted_data:
extracted_cards = extracted_data['cards']
elif 'memory_cards' in extracted_data:
extracted_cards = extracted_data['memory_cards']
else:
# Single card
extracted_cards = [extracted_data] if extracted_data else []
elif isinstance(extracted_data, list):
extracted_cards = extracted_data
else:
extracted_cards = []
# Convert extracted data to AdvancedMemoryCard objects
for card_data in extracted_cards:
try:
# Ensure required fields
if 'backstory' not in card_data:
card_data['backstory'] = f"Information extracted from conversation {conversation_id}"
if 'date_created' not in card_data:
card_data['date_created'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
if 'person' not in card_data:
card_data['person'] = 'User (primary)'
if 'relationship' not in card_data:
card_data['relationship'] = 'primary account holder'
# Extract main fields
category = card_data.pop('category', 'general')
card_key = card_data.pop('card_key', f"{category}_{conversation_id[:8]}")
backstory = card_data.pop('backstory')
date_created = card_data.pop('date_created')
person = card_data.pop('person')
relationship = card_data.pop('relationship')
# Create card with remaining fields as data
card = AdvancedMemoryCard(
category=category,
card_key=card_key,
backstory=backstory,
date_created=date_created,
person=person,
relationship=relationship,
data=card_data # All remaining fields
)
cards.append(card)
except Exception as e:
logger.warning(f"Error creating memory card: {e}")
continue
except Exception as e:
logger.error(f"Error generating memory cards with LLM: {e}")
# Fallback to basic extraction if LLM fails
cards = self._fallback_extraction(chunks, conversation_id, full_text)
return cards
def _fallback_extraction(self, chunks: List[ConversationChunk],
conversation_id: str, full_text: str) -> List[AdvancedMemoryCard]:
"""Fallback extraction method if LLM is unavailable"""
cards = []
# Extract basic information based on keywords
categories_keywords = {
'financial': ['bank', 'account', 'money', 'payment', 'credit', 'loan', 'savings', 'checking'],
'medical': ['doctor', 'medical', 'health', 'appointment', 'prescription', 'hospital', 'clinic'],
'travel': ['flight', 'travel', 'trip', 'vacation', 'hotel', 'booking', 'destination'],
'family': ['family', 'child', 'parent', 'spouse', 'daughter', 'son', 'wife', 'husband'],
'work': ['job', 'work', 'employer', 'salary', 'office', 'meeting', 'project', 'colleague']
}
text_lower = full_text.lower()
for category, keywords in categories_keywords.items():
if any(keyword in text_lower for keyword in keywords):
card = AdvancedMemoryCard(
category=category,
card_key=f"{category}_{conversation_id[:8]}",
backstory=f"Information about {category} topics discussed in conversation",
date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
person="User (primary)",
relationship="primary account holder",
data={
"conversation_id": conversation_id,
"topics": [kw for kw in keywords if kw in text_lower][:3],
"extracted_from": f"{len(chunks)} conversation chunks",
"extraction_method": "fallback_keywords"
}
)
cards.append(card)
break # Only add one card per conversation in fallback mode
return cards
def search_with_context(self,
query: str,
top_k: int = 3,
include_memory_cards: bool = True) -> Dict[str, Any]:
"""
Search both contextual chunks and memory cards.
Args:
query: Search query
top_k: Number of results
include_memory_cards: Whether to search memory cards too
Returns:
Combined search results
"""
results = {
"query": query,
"chunk_results": [],
"memory_card_results": [],
"combined_context": ""
}
# Search contextual chunks via retrieval pipeline or fallback to local search
try:
response = requests.post(
f"{self.retrieval_url}/search",
json={
"query": query,
"mode": "hybrid",
"top_k": max(20, top_k),
"rerank_top_k": top_k,
"skip_reranking": False
},
timeout=5
)
response.raise_for_status()
data = response.json()
search_results = data.get("reranked_results", [])
for item in search_results:
metadata = item.get("metadata", {})
chunk_id = metadata.get("doc_id")
if chunk_id and chunk_id in self.contextual_chunks:
contextual_chunk = self.contextual_chunks[chunk_id]
# Create a basic ConversationChunk for compatibility
from chunker import ConversationChunk
basic_chunk = ConversationChunk(
chunk_id=contextual_chunk.chunk_id,
conversation_id=contextual_chunk.conversation_id,
test_id=contextual_chunk.test_id,
chunk_index=contextual_chunk.chunk_index,
start_round=contextual_chunk.start_round,
end_round=contextual_chunk.end_round,
messages=contextual_chunk.messages,
metadata=contextual_chunk.metadata,
context_before=contextual_chunk.context_before,
context_after=contextual_chunk.context_after,
created_at=contextual_chunk.created_at
)
result = ContextualSearchResult(
chunk_id=chunk_id,
score=float(item.get("rerank_score", 0)),
chunk=basic_chunk,
match_type="hybrid",
context=contextual_chunk.context,
contextual_chunk=contextual_chunk
)
results["chunk_results"].append(result.to_dict())
except Exception as e:
logger.error(f"Error searching via retrieval pipeline: {e}")
logger.info("Falling back to local search...")
# Fallback: Local search through contextual chunks
query_lower = query.lower()
local_results = []
for chunk_id, chunk in self.contextual_chunks.items():
# Search in contextualized text (skip chunks with empty text)
if query_lower and chunk.contextualized_text and query_lower in chunk.contextualized_text.lower():
# Calculate a simple relevance score based on frequency
score = chunk.contextualized_text.lower().count(query_lower) / len(chunk.contextualized_text)
local_results.append((score, chunk_id, chunk))
# Sort by score and take top_k
local_results.sort(key=lambda x: x[0], reverse=True)
for score, chunk_id, contextual_chunk in local_results[:top_k]:
# Create a basic ConversationChunk for compatibility
from chunker import ConversationChunk
basic_chunk = ConversationChunk(
chunk_id=contextual_chunk.chunk_id,
conversation_id=contextual_chunk.conversation_id,
test_id=contextual_chunk.test_id,
chunk_index=contextual_chunk.chunk_index,
start_round=contextual_chunk.start_round,
end_round=contextual_chunk.end_round,
messages=contextual_chunk.messages,
metadata=contextual_chunk.metadata,
context_before=contextual_chunk.context_before,
context_after=contextual_chunk.context_after,
created_at=contextual_chunk.created_at
)
results["chunk_results"].append({
"chunk_id": chunk_id,
"score": score,
"context": contextual_chunk.context,
"conversation_id": contextual_chunk.conversation_id,
"rounds": f"{contextual_chunk.start_round}-{contextual_chunk.end_round}",
"text": contextual_chunk.original_text
})
logger.info(f"Local search returned {len(results['chunk_results'])} results")
# Search memory cards
if include_memory_cards:
card_results = self.memory_manager.search_cards(query)
for memory_id, card in card_results[:top_k]:
results["memory_card_results"].append({
"memory_id": memory_id,
"category": card.category,
"backstory": card.backstory,
"person": card.person,
"data": card.data
})
# Build combined context
context_parts = []
if results["memory_card_results"]:
context_parts.append("=== RELEVANT MEMORY CARDS ===")
for card_result in results["memory_card_results"]:
context_parts.append(f"- {card_result['category']}: {card_result['backstory']}")
if results["chunk_results"]:
context_parts.append("\n=== RELEVANT CONVERSATION CHUNKS ===")
for chunk_result in results["chunk_results"]:
context_parts.append(f"- {chunk_result.get('context', 'No context')}")
results["combined_context"] = "\n".join(context_parts)
return results
def get_agent_context(self, max_memory_cards: int = 10) -> str:
"""
Get the complete context for the agent including memory cards.
Args:
max_memory_cards: Maximum number of memory cards to include
Returns:
Formatted context string
"""
return self.memory_manager.get_context_string(max_cards=max_memory_cards)
def save_index(self, path: Optional[str] = None):
"""Save the index and memory to disk"""
path = path or self.index_config.index_path
# Save contextual chunks
chunks_data = {
chunk_id: chunk.to_dict()
for chunk_id, chunk in self.contextual_chunks.items()
}
with open(f"{path}_contextual_chunks.json", 'w', encoding='utf-8') as f:
json.dump(chunks_data, f, ensure_ascii=False, indent=2)
# Memory cards are automatically saved by the memory manager
# Save statistics
with open(f"{path}_stats.json", 'w', encoding='utf-8') as f:
json.dump(self.stats, f, indent=2)
logger.info(f"Saved index to {path}")
def load_index(self, path: Optional[str] = None):
"""Load the index from disk"""
path = path or self.index_config.index_path
try:
# Load contextual chunks
with open(f"{path}_contextual_chunks.json", 'r', encoding='utf-8') as f:
chunks_data = json.load(f)
self.contextual_chunks = {}
for chunk_id, chunk_dict in chunks_data.items():
# Convert messages
from chunker import ConversationMessage
messages = []
for msg_data in chunk_dict.get('messages', []):
messages.append(ConversationMessage(**msg_data))
# Create contextual chunk
chunk = ContextualConversationChunk(
chunk_id=chunk_dict['chunk_id'],
conversation_id=chunk_dict['conversation_id'],
test_id=chunk_dict['test_id'],
chunk_index=chunk_dict['chunk_index'],
start_round=chunk_dict['start_round'],
end_round=chunk_dict['end_round'],
messages=messages,
original_text=chunk_dict.get('original_text', ''),
context=chunk_dict.get('context', ''),
contextualized_text=chunk_dict.get('contextualized_text', ''),
context_tokens=chunk_dict.get('context_tokens', 0),
generation_time=chunk_dict.get('generation_time', 0),
metadata=chunk_dict.get('metadata', {}),
context_before=chunk_dict.get('context_before'),
context_after=chunk_dict.get('context_after'),
created_at=chunk_dict.get('created_at', '')
)
self.contextual_chunks[chunk_id] = chunk
# Memory cards are automatically loaded by the memory manager
# Load statistics
stats_path = f"{path}_stats.json"
if Path(stats_path).exists():
with open(stats_path, 'r', encoding='utf-8') as f:
self.stats = json.load(f)
logger.info(f"Loaded {len(self.contextual_chunks)} contextual chunks from {path}")
# Re-index in retrieval pipeline
self._index_contextual_chunks(list(self.contextual_chunks.values()))
except Exception as e:
logger.error(f"Error loading index: {e}")
raise
def get_statistics(self) -> Dict[str, Any]:
"""Get comprehensive statistics"""
stats = self.stats.copy()
# Add chunker statistics
stats["chunker_stats"] = self.contextual_chunker.get_statistics()
# Add memory manager statistics
stats["memory_stats"] = self.memory_manager.get_statistics()
# Calculate averages
if stats["chunks_indexed"] > 0:
stats["avg_indexing_time"] = stats["indexing_time"] / stats["chunks_indexed"]
return stats
demo_agent_logging.py¶
#!/usr/bin/env python3
"""Demonstration of agent tool logging with full content"""
import os
import json
import logging
from rich.console import Console
from chunker import ConversationChunker, ConversationChunk, ConversationMessage
from indexer import MemoryIndexer
from agent import UserMemoryRAGAgent
from config import Config
# Set up logging to see agent logs
logging.basicConfig(
level=logging.INFO,
format='%(message)s' # Simple format to show just the message
)
# Set dummy API key for demo
os.environ["KIMI_API_KEY"] = "test_key"
console = Console()
def create_demo_conversation():
"""Create a demo conversation with important information"""
messages = [
ConversationMessage(role="user", content="Hi, I'd like to open a checking account."),
ConversationMessage(role="assistant", content="I'd be happy to help you open a checking account. May I have your full name?"),
ConversationMessage(role="user", content="John Michael Smith"),
ConversationMessage(role="assistant", content="Thank you, Mr. Smith. What's your date of birth?"),
ConversationMessage(role="user", content="March 15, 1985"),
ConversationMessage(role="assistant", content="Perfect. And your Social Security Number?"),
ConversationMessage(role="user", content="123-45-6789"),
ConversationMessage(role="assistant", content="Thank you. What's your current address?"),
ConversationMessage(role="user", content="456 Oak Avenue, Springfield, IL 62701"),
ConversationMessage(role="assistant", content="Great. And a phone number where we can reach you?"),
ConversationMessage(role="user", content="555-0123"),
ConversationMessage(role="assistant", content="Excellent. I've set up your account. Your new account number is 4429853327."),
ConversationMessage(role="user", content="Great! What's the routing number?"),
ConversationMessage(role="assistant", content="The routing number for our Springfield branch is 071000013."),
ConversationMessage(role="user", content="Perfect, thank you for your help!"),
ConversationMessage(role="assistant", content="You're welcome! Your debit card will arrive in 7-10 business days."),
]
chunk = ConversationChunk(
chunk_id="demo_bank_account_001",
conversation_id="bank_conv_001",
test_id="demo_test",
chunk_index=0,
start_round=1,
end_round=8,
messages=messages,
metadata={
"business": "First National Bank",
"department": "New Accounts",
"date": "2024-11-15",
"agent": "Sarah Johnson"
}
)
return [chunk]
def demonstrate_agent_logging():
"""Demonstrate how the agent logs tool calls and results"""
console.print("\n[bold cyan]Agent Tool Logging Demonstration[/bold cyan]")
console.print("="*80)
console.print("[yellow]Note: Watch for tool call parameters and full results in the logs[/yellow]\n")
# Initialize configuration
config = Config.from_env()
config.agent.enable_reasoning = True # Enable reasoning for more detailed logs
# Initialize indexer and add demo conversation
indexer = MemoryIndexer(config.index)
chunks = create_demo_conversation()
# Manually add chunks without pipeline indexing for demo
for chunk in chunks:
indexer.chunks[chunk.chunk_id] = chunk
indexer.chunk_texts[chunk.chunk_id] = chunk.to_text()
# Initialize agent
agent = UserMemoryRAGAgent(indexer, config)
# Test question that will trigger tool usage
test_question = "What is the customer's account number and when will they receive their debit card?"
console.print(f"\n[bold]Test Question:[/bold] {test_question}")
console.print("\n[yellow]Agent processing (watch the tool logs below):[/yellow]\n")
console.print("="*80 + "\n")
# This would normally call the LLM, but for demo we'll simulate tool calls
# In a real scenario, the agent.answer_question() method would be called
# Simulate what the agent would do:
from tools import MemoryTools
tools = MemoryTools(indexer)
# Simulate search_memory call (this is what the agent would do)
console.print("[dim]Agent would execute:[/dim]")
result1 = agent._execute_tool("search_memory", {
"query": "account number debit card",
"top_k": 3,
"filter_test_id": "demo_test"
})
# The logging happens automatically in _execute_tool
console.print("\n[dim]Agent might also execute:[/dim]")
result2 = agent._execute_tool("get_full_conversation", {
"conversation_id": "bank_conv_001",
"test_id": "demo_test"
})
console.print("\n" + "="*80)
console.print("[green]Demonstration complete![/green]")
console.print("\nKey observations:")
console.print(" 1. Tool call parameters are logged in full")
console.print(" 2. Tool results show complete content (not truncated)")
console.print(" 3. Each tool call is clearly separated with dividers")
console.print(" 4. JSON formatting makes results easy to read")
def main():
"""Run the demonstration"""
console.print("\n[bold]Agentic RAG - Tool Logging Demonstration[/bold]")
console.print("This shows how tool calls and results are logged to the console\n")
try:
demonstrate_agent_logging()
except Exception as e:
console.print(f"\n[red]Error during demonstration: {e}[/red]")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()
demo_ui_improvements.py¶
#!/usr/bin/env python3
"""Demonstration of UI improvements for test case selection"""
import os
import sys
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
# Mock environment for demo
os.environ["KIMI_API_KEY"] = "demo_key"
console = Console()
def demo_single_test_selection():
"""Demonstrate the improved single test case selection"""
console.print("\n[bold cyan]Improved Single Test Case Selection[/bold cyan]")
console.print("="*60)
# Simulated test cases
test_cases = {
"layer1_01_bank_account": "Bank Account Setup - Personal Details Retrieval",
"layer1_02_insurance_claim": "Insurance Claim - Policy and Representative Information",
"layer1_03_medical_appointment": "Medical Appointment - Doctor and Schedule Details",
"layer2_01_multiple_vehicles": "Multiple Vehicles - Service Scheduling Conflict",
"layer2_02_multiple_properties": "Multiple Properties - Maintenance Priorities",
"layer3_01_travel_coordination": "Travel Coordination - Multi-Factor Planning"
}
console.print("\n[bold]Available Test Cases:[/bold]")
# Group by category
categories = {
"layer1": [],
"layer2": [],
"layer3": []
}
for test_id, title in test_cases.items():
cat = test_id.split("_")[0]
categories[cat].append((test_id, title))
# Display with numbers
test_ids_list = []
for cat in sorted(categories.keys()):
console.print(f"\n[cyan]{cat.upper()}:[/cyan]")
for test_id, title in categories[cat]:
test_ids_list.append(test_id)
idx = len(test_ids_list)
console.print(f" [{idx}] {test_id}: {title[:60]}...")
console.print("\n[dim]Enter test ID directly or number from the list above[/dim]")
console.print("[yellow]Example inputs: '3' or 'layer1_03_medical_appointment'[/yellow]")
def demo_view_test_cases_table():
"""Demonstrate the improved test cases table with index numbers"""
console.print("\n[bold cyan]Improved Test Cases Table View[/bold cyan]")
console.print("="*60)
# Create table with index numbers
table = Table(title="Loaded Test Cases")
table.add_column("#", style="dim", width=4)
table.add_column("ID", style="cyan")
table.add_column("Category", style="magenta")
table.add_column("Title", style="green")
table.add_column("Conversations", justify="right")
# Sample data
test_data = [
("layer1_01_bank_account", "layer1", "Bank Account Setup - Personal Details Retrieval", "1"),
("layer1_02_insurance_claim", "layer1", "Insurance Claim - Policy and Representative Info", "1"),
("layer2_01_multiple_vehicles", "layer2", "Multiple Vehicles - Service Scheduling Conflict", "2"),
("layer3_01_travel_coordination", "layer3", "Travel Coordination - Multi-Factor Planning", "3")
]
for idx, (test_id, category, title, convs) in enumerate(test_data, 1):
table.add_row(
str(idx),
test_id,
category,
title[:50] + "..." if len(title) > 50 else title,
convs
)
console.print(table)
console.print("\n[yellow]Users can now select by entering either:[/yellow]")
console.print(" • The index number (e.g., '2')")
console.print(" • The full test ID (e.g., 'layer1_02_insurance_claim')")
def demo_category_evaluation():
"""Demonstrate the improved category evaluation selection"""
console.print("\n[bold cyan]Improved Category Evaluation[/bold cyan]")
console.print("="*60)
console.print("\n[bold]Available Categories:[/bold]")
console.print(" layer1: 20 test cases")
console.print(" layer2: 20 test cases")
console.print(" layer3: 20 test cases")
console.print("\n[dim]After selecting a category, you'll see:[/dim]")
console.print(f"\n[cyan]Test cases in layer1:[/cyan]")
sample_cases = [
("layer1_01_bank_account", "Bank Account Setup - Personal Details Retrieval"),
("layer1_02_insurance_claim", "Insurance Claim - Policy and Representative Info"),
("layer1_03_medical_appointment", "Medical Appointment - Doctor and Schedule Details")
]
for test_id, title in sample_cases:
console.print(f" • {test_id}: {title[:60]}...")
console.print(" [dim]... (17 more test cases)[/dim]")
console.print(f"\n[cyan]Total: 20 test cases[/cyan]")
console.print("[yellow]User can review the list before confirming evaluation[/yellow]")
def main():
"""Run all UI improvement demonstrations"""
console.print(Panel.fit(
"[bold]UI Improvements Demo - Test Case Selection[/bold]\n"
"Showing how the improved interface helps users select test cases",
border_style="cyan"
))
# Demo 1: Single test selection
demo_single_test_selection()
# Demo 2: Table view with indices
console.print("\n" + "="*80 + "\n")
demo_view_test_cases_table()
# Demo 3: Category evaluation
console.print("\n" + "="*80 + "\n")
demo_category_evaluation()
console.print("\n" + "="*80)
console.print("[bold green]Summary of Improvements:[/bold green]")
console.print("✓ Test cases are displayed with index numbers for easy selection")
console.print("✓ Users can select by number OR by test ID")
console.print("✓ Categories show test case counts before evaluation")
console.print("✓ Test case lists are shown before batch evaluation")
console.print("✓ No need to memorize or guess test IDs")
console.print("="*80 + "\n")
if __name__ == "__main__":
main()
evaluator.py¶
"""Evaluation Framework Integration for User Memory RAG Agent
This module integrates with the user-memory-evaluation framework to load
test cases and evaluate the agent's performance.
"""
import os
import sys
import json
import yaml
import logging
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
from dataclasses import dataclass, field
from datetime import datetime
import openai
# Import our local modules first to avoid conflicts
from config import Config
from chunker import ConversationChunker
from indexer import MemoryIndexer
from agent import UserMemoryRAGAgent
# Import modules from user-memory-evaluation project
LLMEvaluator = None
EvalTestCase = None
ConversationHistory = None
EvalMessage = None
MessageRole = None
# Import the specific modules we need from user-memory-evaluation
eval_project_path = Path(__file__).parent.parent.parent / "week2" / "user-memory-evaluation"
if eval_project_path.exists():
import importlib.util
try:
# Load models module from user-memory-evaluation
models_spec = importlib.util.spec_from_file_location(
"user_memory_models",
eval_project_path / "models.py"
)
models_module = importlib.util.module_from_spec(models_spec)
models_spec.loader.exec_module(models_module)
# Load config module from user-memory-evaluation for the evaluator
eval_config_spec = importlib.util.spec_from_file_location(
"user_memory_config",
eval_project_path / "config.py"
)
eval_config_module = importlib.util.module_from_spec(eval_config_spec)
eval_config_spec.loader.exec_module(eval_config_module)
# Load evaluator module from user-memory-evaluation
eval_spec = importlib.util.spec_from_file_location(
"user_memory_evaluator",
eval_project_path / "evaluator.py"
)
eval_module = importlib.util.module_from_spec(eval_spec)
# Temporarily add required modules to sys.modules for the evaluator to find
sys.modules['models'] = models_module
sys.modules['config'] = eval_config_module
# Execute the evaluator module
eval_spec.loader.exec_module(eval_module)
# Clean up sys.modules to avoid conflicts
del sys.modules['models']
sys.modules['config'] = sys.modules[Config.__module__] # Restore our own config
# Extract the classes we need
LLMEvaluator = eval_module.LLMEvaluator
EvalTestCase = models_module.TestCase
ConversationHistory = models_module.ConversationHistory
EvalMessage = models_module.ConversationMessage
MessageRole = models_module.MessageRole
print(f"Successfully imported LLMEvaluator from {eval_project_path}")
except Exception as e:
print(f"Error loading LLMEvaluator: {e}")
import traceback
traceback.print_exc()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class TestCase:
"""Test case matching user-memory-evaluation framework structure"""
test_id: str
category: str
title: str
description: str
conversation_histories: List[Dict[str, Any]]
user_question: str
evaluation_criteria: str # Primary evaluation criteria (was expected_answer)
expected_behavior: Optional[str] = None # Optional expected behavior
metadata: Dict[str, Any] = field(default_factory=dict)
@dataclass
class EvaluationResult:
"""Result from evaluating a test case"""
test_id: str
success: bool
agent_answer: str
evaluation_criteria: str # Changed from expected_answer to match TestCase
iterations: int
tool_calls: int
trajectory: Optional[Dict[str, Any]] = None
processing_time: float = 0.0
indexing_time: float = 0.0
chunk_count: int = 0
llm_evaluation: Optional[Dict[str, Any]] = None # LLM evaluation details
class UserMemoryEvaluator:
"""Evaluates the RAG agent on user memory test cases"""
def __init__(self, config: Optional[Config] = None):
"""
Initialize the evaluator
Args:
config: Configuration object
"""
self.config = config or Config.from_env()
self.test_cases: Dict[str, TestCase] = {}
self.results: Dict[str, EvaluationResult] = {}
# Components
self.chunker = ConversationChunker(self.config.chunking)
self.indexer = None
self.agent = None
# Initialize LLM evaluator if available
self.llm_evaluator = None
if LLMEvaluator:
try:
self.llm_evaluator = LLMEvaluator()
logger.info("LLM Evaluator initialized for automatic evaluation")
except Exception as e:
logger.warning(f"Could not initialize LLM evaluator: {e}")
logger.info("Automatic LLM evaluation will be skipped")
else:
logger.info("LLM Evaluator not available - automatic evaluation will be skipped")
# Paths
self.test_cases_dir = Path(self.config.evaluation.test_cases_dir)
self.results_dir = Path(self.config.evaluation.results_dir)
self.results_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Initialized evaluator with test cases from: {self.test_cases_dir}")
def load_test_cases(self, category: Optional[str] = None) -> List[str]:
"""
Load test cases from YAML files
Args:
category: Optional category filter (layer1, layer2, layer3)
Returns:
List of test case IDs that were loaded
"""
test_case_ids = []
# Determine which categories to load
if category:
categories = [category]
else:
categories = ["layer1", "layer2", "layer3"]
for cat in categories:
category_dir = self.test_cases_dir / cat
if not category_dir.exists():
logger.warning(f"Category directory {category_dir} does not exist")
continue
# Load all YAML files in category
for yaml_file in category_dir.glob("*.yaml"):
try:
test_case = self._load_single_test_case(yaml_file)
if test_case:
test_case_ids.append(test_case.test_id)
self.test_cases[test_case.test_id] = test_case
except Exception as e:
logger.error(f"Error loading {yaml_file}: {e}")
logger.info(f"Loaded {len(test_case_ids)} test cases")
return test_case_ids
def _load_single_test_case(self, yaml_file: Path) -> Optional[TestCase]:
"""Load a single test case from YAML file"""
with open(yaml_file, 'r', encoding='utf-8') as f:
data = yaml.safe_load(f)
if not data:
return None
# Parse conversation histories
conversation_histories = []
for conv_data in data.get('conversation_histories', []):
# Ensure messages are in the right format
messages = []
msg_list = conv_data.get('messages', [])
for msg in msg_list:
if isinstance(msg, dict) and 'role' in msg and 'content' in msg:
messages.append(msg)
conversation = {
"conversation_id": conv_data.get('conversation_id', ''),
"timestamp": conv_data.get('timestamp', ''),
"metadata": conv_data.get('metadata', {}),
"messages": messages
}
conversation_histories.append(conversation)
# Load evaluation criteria and expected behavior directly from YAML
evaluation_criteria = data.get('evaluation_criteria', '')
expected_behavior = data.get('expected_behavior', None)
# Create test case with matching structure
test_case = TestCase(
test_id=data.get('test_id', ''),
category=data.get('category', ''),
title=data.get('title', ''),
description=data.get('description', ''),
conversation_histories=conversation_histories,
user_question=data.get('user_question', ''),
evaluation_criteria=evaluation_criteria,
expected_behavior=expected_behavior,
metadata=data.get('metadata', {})
)
return test_case
def prepare_test_case(self, test_id: str) -> Tuple[int, float]:
"""
Prepare a test case by chunking and indexing its conversations
Args:
test_id: The test case ID to prepare
Returns:
Tuple of (number of chunks, indexing time)
"""
if test_id not in self.test_cases:
raise ValueError(f"Test case {test_id} not found")
test_case = self.test_cases[test_id]
start_time = datetime.now()
logger.info(f"Preparing test case: {test_id}")
# Create new indexer for this test case
self.indexer = MemoryIndexer(self.config.index)
# Chunk all conversations
all_chunks = []
for conv_history in test_case.conversation_histories:
chunks = self.chunker.chunk_conversation(
conversation_id=conv_history['conversation_id'],
test_id=test_id,
messages=conv_history['messages'],
metadata=conv_history.get('metadata', {})
)
all_chunks.extend(chunks)
logger.info(f"Created {len(all_chunks)} chunks for test case {test_id}")
# Index chunks
self.indexer.add_chunks(all_chunks)
# Save index if caching is enabled
if self.config.evaluation.enable_caching:
cache_path = self.results_dir / f"index_{test_id}"
self.indexer.save_index(str(cache_path))
logger.info(f"Cached index for {test_id}")
# Create agent with the indexed data
self.agent = UserMemoryRAGAgent(self.indexer, self.config)
end_time = datetime.now()
indexing_time = (end_time - start_time).total_seconds()
return len(all_chunks), indexing_time
def evaluate_test_case(self, test_id: str) -> EvaluationResult:
"""
Evaluate a single test case
Args:
test_id: The test case ID to evaluate
Returns:
Evaluation result
"""
if test_id not in self.test_cases:
raise ValueError(f"Test case {test_id} not found")
test_case = self.test_cases[test_id]
logger.info(f"\n{'='*60}")
logger.info(f"Evaluating: {test_case.title}")
logger.info(f"Category: {test_case.category}")
logger.info(f"Question: {test_case.user_question}")
logger.info(f"{'='*60}")
# Check if we can load cached index
chunk_count = 0
indexing_time = 0.0
if self.config.evaluation.enable_caching:
cache_path = self.results_dir / f"index_{test_id}"
if Path(f"{cache_path}_chunks.json").exists():
logger.info(f"Loading cached index for {test_id}")
self.indexer = MemoryIndexer(self.config.index)
self.indexer.load_index(str(cache_path))
self.agent = UserMemoryRAGAgent(self.indexer, self.config)
chunk_count = len(self.indexer.chunks)
else:
chunk_count, indexing_time = self.prepare_test_case(test_id)
else:
chunk_count, indexing_time = self.prepare_test_case(test_id)
# Get agent's answer
start_time = datetime.now()
result = self.agent.answer_question(
question=test_case.user_question,
test_id=test_id,
stream=False
)
end_time = datetime.now()
processing_time = (end_time - start_time).total_seconds()
agent_answer = result.get("answer", "")
# Perform LLM evaluation if available
llm_evaluation = None
if self.llm_evaluator and agent_answer:
logger.info("\n" + "="*60)
logger.info("Running LLM Evaluation...")
logger.info("-"*60)
try:
# Convert test case to format expected by LLM evaluator
# Using the imported classes from user-memory-evaluation models
# Build conversation histories for evaluator
eval_histories = []
for conv_hist in test_case.conversation_histories:
eval_messages = []
for msg in conv_hist.get('messages', []):
eval_messages.append(EvalMessage(
role=MessageRole(msg.get('role', 'user')),
content=msg.get('content', '')
))
# Extract timestamp from metadata or use a default
timestamp = conv_hist.get('timestamp', '')
if not timestamp and 'metadata' in conv_hist:
# Try to extract from metadata
metadata = conv_hist.get('metadata', {})
timestamp = metadata.get('timestamp', metadata.get('date', '2024-01-01'))
if not timestamp:
timestamp = '2024-01-01' # Default timestamp
eval_histories.append(ConversationHistory(
conversation_id=conv_hist.get('conversation_id', ''),
timestamp=timestamp,
messages=eval_messages,
metadata=conv_hist.get('metadata', {})
))
# Use the test case's evaluation_criteria directly
eval_test_case = EvalTestCase(
test_id=test_case.test_id,
category=test_case.category,
title=test_case.title,
description=test_case.description,
conversation_histories=eval_histories,
user_question=test_case.user_question,
evaluation_criteria=test_case.evaluation_criteria if test_case.evaluation_criteria else "The agent should provide a relevant and accurate response based on the conversation history.",
expected_behavior=test_case.expected_behavior # Pass through expected_behavior
)
# Run LLM evaluation
llm_result = self.llm_evaluator.evaluate(
test_case=eval_test_case,
agent_response=agent_answer,
extracted_memory=None
)
llm_evaluation = {
"reward": llm_result.reward,
"passed": llm_result.passed if llm_result.passed is not None else llm_result.reward >= 0.6,
"reasoning": llm_result.reasoning,
"required_info_found": llm_result.required_info_found if hasattr(llm_result, 'required_info_found') else {},
"suggestions": llm_result.suggestions if hasattr(llm_result, 'suggestions') else None
}
# Log evaluation results with full reasoning
logger.info("-"*60)
logger.info(f"LLM Evaluation Reward: {llm_result.reward:.3f}/1.000")
logger.info(f"Passed: {'Yes' if llm_evaluation['passed'] else 'No'}")
logger.info("-"*60)
logger.info(f"Evaluation Reasoning:")
logger.info(llm_result.reasoning)
logger.info("-"*60)
if llm_result.required_info_found:
logger.info("Required Information Found:")
for info, found in llm_result.required_info_found.items():
check = "✓" if found else "✗"
logger.info(f" {check} {info}")
if llm_result.suggestions:
logger.info(f"Suggestions: {llm_result.suggestions}")
logger.info("="*60)
except Exception as e:
logger.error(f"Error during LLM evaluation: {e}")
llm_evaluation = {"error": str(e)}
# Create evaluation result
eval_result = EvaluationResult(
test_id=test_id,
success=llm_evaluation.get('passed', result.get("success", False)) if llm_evaluation else result.get("success", False),
agent_answer=agent_answer,
evaluation_criteria=test_case.evaluation_criteria, # Use evaluation_criteria
iterations=result.get("iterations", 0),
tool_calls=result.get("tool_calls", 0),
trajectory=result.get("trajectory"),
processing_time=processing_time,
indexing_time=indexing_time,
chunk_count=chunk_count
)
# Add LLM evaluation details to the result if available
if llm_evaluation:
eval_result.llm_evaluation = llm_evaluation
# Save result
self.results[test_id] = eval_result
# Save trajectory if enabled
if self.config.evaluation.save_trajectories and eval_result.trajectory:
trajectory_file = self.results_dir / f"trajectory_{test_id}.json"
with open(trajectory_file, 'w', encoding='utf-8') as f:
json.dump(eval_result.trajectory, f, ensure_ascii=False, indent=2)
# Log summary
logger.info(f"\n{'='*60}")
logger.info(f"Evaluation Complete for {test_id}")
if llm_evaluation and 'reward' in llm_evaluation:
logger.info(f"LLM Evaluation Passed: {'✓' if eval_result.success else '✗'}")
logger.info(f"LLM Reward Score: {llm_evaluation['reward']:.3f}/1.000")
else:
logger.info(f"Success: {eval_result.success}")
logger.info(f"Iterations: {eval_result.iterations}")
logger.info(f"Tool Calls: {eval_result.tool_calls}")
logger.info(f"Chunks: {eval_result.chunk_count}")
logger.info(f"Processing Time: {eval_result.processing_time:.2f}s")
logger.info(f"Indexing Time: {eval_result.indexing_time:.2f}s")
logger.info(f"{'='*60}")
return eval_result
def evaluate_batch(self,
test_ids: Optional[List[str]] = None,
category: Optional[str] = None) -> Dict[str, EvaluationResult]:
"""
Evaluate multiple test cases
Args:
test_ids: List of test IDs to evaluate (evaluates all if None)
category: Category filter if test_ids not provided
Returns:
Dictionary of results
"""
# Determine which test cases to evaluate
if test_ids:
ids_to_evaluate = test_ids
else:
# Load test cases if needed
if not self.test_cases:
self.load_test_cases(category)
if category:
ids_to_evaluate = [
tid for tid, tc in self.test_cases.items()
if tc.category == category
]
else:
ids_to_evaluate = list(self.test_cases.keys())
logger.info(f"Evaluating {len(ids_to_evaluate)} test cases")
# Evaluate each test case
for i, test_id in enumerate(ids_to_evaluate, 1):
logger.info(f"\n[{i}/{len(ids_to_evaluate)}] Evaluating {test_id}")
try:
self.evaluate_test_case(test_id)
except Exception as e:
logger.error(f"Error evaluating {test_id}: {e}")
self.results[test_id] = EvaluationResult(
test_id=test_id,
success=False,
agent_answer=f"Error: {str(e)}",
evaluation_criteria="",
iterations=0,
tool_calls=0
)
return self.results
def generate_report(self, output_file: Optional[str] = None) -> str:
"""
Generate evaluation report
Args:
output_file: Optional file path to save report
Returns:
Report as string
"""
if not self.results:
return "No evaluation results available"
lines = []
lines.append("="*80)
lines.append("USER MEMORY RAG EVALUATION REPORT")
lines.append(f"Generated: {datetime.now().isoformat()}")
lines.append("="*80)
lines.append("")
# Summary statistics
total = len(self.results)
successful = sum(1 for r in self.results.values() if r.success)
# Calculate LLM evaluation metrics if available
llm_evaluated = sum(1 for r in self.results.values() if r.llm_evaluation and 'reward' in r.llm_evaluation)
avg_reward = 0.0
if llm_evaluated > 0:
avg_reward = sum(r.llm_evaluation['reward'] for r in self.results.values()
if r.llm_evaluation and 'reward' in r.llm_evaluation) / llm_evaluated
lines.append("SUMMARY")
lines.append("-"*40)
lines.append(f"Total Test Cases: {total}")
lines.append(f"Successful: {successful}/{total} ({100*successful/total:.1f}%)")
if llm_evaluated > 0:
lines.append(f"LLM Evaluated: {llm_evaluated}/{total}")
lines.append(f"Average LLM Reward: {avg_reward:.3f}/1.000")
lines.append("")
# Average metrics
avg_iterations = sum(r.iterations for r in self.results.values()) / total
avg_tool_calls = sum(r.tool_calls for r in self.results.values()) / total
avg_chunks = sum(r.chunk_count for r in self.results.values()) / total
avg_proc_time = sum(r.processing_time for r in self.results.values()) / total
avg_idx_time = sum(r.indexing_time for r in self.results.values()) / total
lines.append("AVERAGE METRICS")
lines.append("-"*40)
lines.append(f"Iterations per test: {avg_iterations:.2f}")
lines.append(f"Tool calls per test: {avg_tool_calls:.2f}")
lines.append(f"Chunks per test: {avg_chunks:.1f}")
lines.append(f"Processing time: {avg_proc_time:.2f}s")
lines.append(f"Indexing time: {avg_idx_time:.2f}s")
lines.append("")
# Results by category
categories = {}
for test_id, result in self.results.items():
if test_id in self.test_cases:
cat = self.test_cases[test_id].category
if cat not in categories:
categories[cat] = {"total": 0, "successful": 0}
categories[cat]["total"] += 1
if result.success:
categories[cat]["successful"] += 1
lines.append("RESULTS BY CATEGORY")
lines.append("-"*40)
for cat, stats in sorted(categories.items()):
pct = 100 * stats["successful"] / stats["total"]
lines.append(f"{cat}: {stats['successful']}/{stats['total']} ({pct:.1f}%)")
lines.append("")
# Individual test results
lines.append("INDIVIDUAL TEST RESULTS")
lines.append("-"*40)
for test_id, result in sorted(self.results.items()):
status = "✓" if result.success else "✗"
test_title = self.test_cases.get(test_id, {}).title if test_id in self.test_cases else test_id
lines.append(f"{status} {test_id}: {test_title}")
lines.append(f" Iterations: {result.iterations}, Tool calls: {result.tool_calls}")
lines.append(f" Processing: {result.processing_time:.2f}s, Chunks: {result.chunk_count}")
if result.llm_evaluation and 'reward' in result.llm_evaluation:
lines.append(f" LLM Reward: {result.llm_evaluation['reward']:.3f}/1.000")
lines.append("")
report = "\n".join(lines)
# Save report if output file provided
if output_file:
with open(output_file, 'w', encoding='utf-8') as f:
f.write(report)
logger.info(f"Report saved to {output_file}")
return report
def save_results(self, output_file: Optional[str] = None):
"""
Save evaluation results to JSON
Args:
output_file: Output file path (defaults to timestamped file)
"""
if not output_file:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = self.results_dir / f"results_{timestamp}.json"
results_data = {}
for test_id, result in self.results.items():
results_data[test_id] = {
"success": result.success,
"agent_answer": result.agent_answer,
"evaluation_criteria": result.evaluation_criteria,
"iterations": result.iterations,
"tool_calls": result.tool_calls,
"processing_time": result.processing_time,
"indexing_time": result.indexing_time,
"chunk_count": result.chunk_count
}
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(results_data, f, ensure_ascii=False, indent=2)
logger.info(f"Results saved to {output_file}")
indexer.py¶
"""RAG Indexer for User Memory Conversations
This module handles indexing of conversation chunks using the retrieval pipeline service.
Interfaces with the existing retrieval pipeline on port 4242.
"""
import os
import json
import logging
import requests
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
from pathlib import Path
from config import IndexConfig, IndexMode
from chunker import ConversationChunk
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
@dataclass
class SearchResult:
"""Result from searching the index"""
chunk_id: str
score: float
chunk: ConversationChunk
match_type: str # "dense", "sparse", or "hybrid"
def to_dict(self) -> Dict[str, Any]:
return {
"chunk_id": self.chunk_id,
"score": self.score,
"match_type": self.match_type,
"conversation_id": self.chunk.conversation_id,
"test_id": self.chunk.test_id,
"rounds": f"{self.chunk.start_round}-{self.chunk.end_round}",
"text": self.chunk.to_text()
}
class MemoryIndexer:
"""Indexes and searches conversation chunks using the retrieval pipeline service"""
def __init__(self, config: Optional[IndexConfig] = None):
"""
Initialize the indexer
Args:
config: Index configuration
"""
self.config = config or IndexConfig()
self.chunks: Dict[str, ConversationChunk] = {}
self.chunk_texts: Dict[str, str] = {} # Map chunk_id to prepared text
self.doc_id_mapping: Dict[str, str] = {} # Map generated doc_id to our chunk_id
# Retrieval pipeline URL
self.retrieval_url = "http://localhost:4242"
# Create directories
Path(self.config.index_path).parent.mkdir(parents=True, exist_ok=True)
Path(self.config.chunk_store_path).parent.mkdir(parents=True, exist_ok=True)
# Check retrieval pipeline availability
self._check_retrieval_pipeline()
logger.info(f"Initialized indexer with mode: {self.config.mode}")
def _check_retrieval_pipeline(self):
"""Check if the retrieval pipeline service is available"""
try:
response = requests.get(f"{self.retrieval_url}/health", timeout=2)
if response.status_code == 200:
logger.info("✓ Retrieval pipeline service is available")
else:
logger.warning(f"Retrieval pipeline returned status {response.status_code}")
except requests.exceptions.RequestException as e:
logger.warning(f"Retrieval pipeline service not available at {self.retrieval_url}: {e}")
logger.info("Note: The retrieval pipeline needs to be running. Start it with:")
logger.info(" cd projects/week3/retrieval-pipeline && python api_server.py")
def add_chunks(self, chunks: List[ConversationChunk], rebuild: bool = True):
"""
Add conversation chunks to the index
Args:
chunks: List of conversation chunks to index
rebuild: Whether to rebuild indexes after adding (for retrieval pipeline)
"""
documents = []
for chunk in chunks:
chunk_id = chunk.chunk_id
# Store chunk locally
self.chunks[chunk_id] = chunk
# Prepare text for indexing
chunk_text = self._prepare_chunk_text(chunk)
self.chunk_texts[chunk_id] = chunk_text
# Prepare document for retrieval pipeline
doc = {
"text": chunk_text,
"metadata": {
"doc_id": chunk_id,
"test_id": chunk.test_id,
"conversation_id": chunk.conversation_id,
"chunk_index": chunk.chunk_index,
"start_round": chunk.start_round,
"end_round": chunk.end_round,
**chunk.metadata
}
}
documents.append(doc)
logger.debug(f"Added chunk {chunk_id} to index")
if rebuild and documents:
self._index_documents(documents)
logger.info(f"Added {len(chunks)} chunks to index. Total chunks: {len(self.chunks)}")
def _prepare_chunk_text(self, chunk: ConversationChunk) -> str:
"""
Prepare chunk text for indexing with contextual enrichment
Args:
chunk: Conversation chunk
Returns:
Enriched text for indexing
"""
if not self.config.enable_contextual:
return chunk.to_text()
# Build enriched text with contextual information
lines = []
# Add test case context
lines.append(f"Test Case: {chunk.test_id}")
lines.append(f"Conversation: {chunk.conversation_id}")
# Add metadata as searchable text
if chunk.metadata:
for key, value in chunk.metadata.items():
lines.append(f"{key}: {value}")
# Add the main chunk content
lines.append(chunk.to_text())
# Add semantic tags for better retrieval
lines.append(self._generate_semantic_tags(chunk))
return "\n".join(lines)
def _generate_semantic_tags(self, chunk: ConversationChunk) -> str:
"""
Generate semantic tags for better retrieval
Args:
chunk: Conversation chunk
Returns:
Semantic tags as string
"""
tags = []
# Analyze content for common topics
content = chunk.to_text().lower()
# Financial topics
if any(word in content for word in ["account", "bank", "credit", "loan", "payment"]):
tags.append("financial")
# Insurance topics
if any(word in content for word in ["insurance", "claim", "policy", "coverage"]):
tags.append("insurance")
# Medical topics
if any(word in content for word in ["medical", "doctor", "appointment", "prescription"]):
tags.append("medical")
# Travel topics
if any(word in content for word in ["flight", "hotel", "travel", "booking", "reservation"]):
tags.append("travel")
# Add position tags
if chunk.chunk_index == 0:
tags.append("conversation_start")
# Add round count tags
round_count = chunk.end_round - chunk.start_round + 1
if round_count < 10:
tags.append("short_segment")
elif round_count > 30:
tags.append("long_segment")
return f"Tags: {', '.join(tags)}" if tags else ""
def _index_documents(self, documents: List[Dict[str, Any]]):
"""Send documents to the retrieval pipeline for indexing"""
try:
# First, clear existing index
clear_response = requests.post(f"{self.retrieval_url}/clear")
if clear_response.status_code == 200:
logger.info("Cleared existing index")
# Index documents one by one (retrieval pipeline expects individual documents)
indexed_count = 0
failed_count = 0
for doc in documents:
try:
response = requests.post(
f"{self.retrieval_url}/index",
json=doc # Send individual document
)
if response.status_code == 200:
result = response.json()
generated_doc_id = result.get("doc_id")
our_chunk_id = doc.get("metadata", {}).get("doc_id")
# Store the mapping between generated doc_id and our chunk_id
if generated_doc_id and our_chunk_id:
self.doc_id_mapping[generated_doc_id] = our_chunk_id
indexed_count += 1
else:
failed_count += 1
logger.warning(f"Failed to index document: {doc.get('metadata', {}).get('doc_id', 'unknown')}")
except requests.exceptions.RequestException as e:
failed_count += 1
logger.warning(f"Error indexing document: {e}")
logger.info(f"Indexed {indexed_count} documents successfully ({failed_count} failed)")
except requests.exceptions.RequestException as e:
logger.error(f"Error connecting to retrieval pipeline: {e}")
logger.info("Make sure the retrieval pipeline is running on port 4242")
def build_indexes(self):
"""Build or rebuild indexes by sending all chunks to retrieval pipeline"""
if not self.chunks:
logger.warning("No chunks to index")
return
# Prepare all documents
documents = []
for chunk_id, chunk in self.chunks.items():
chunk_text = self.chunk_texts.get(chunk_id) or self._prepare_chunk_text(chunk)
doc = {
"text": chunk_text,
"metadata": {
"doc_id": chunk_id,
"test_id": chunk.test_id,
"conversation_id": chunk.conversation_id,
"chunk_index": chunk.chunk_index,
"start_round": chunk.start_round,
"end_round": chunk.end_round,
**chunk.metadata
}
}
documents.append(doc)
# Send to retrieval pipeline
self._index_documents(documents)
logger.info("Index building complete")
def search(self,
query: str,
top_k: int = 3,
mode: Optional[IndexMode] = None) -> List[SearchResult]:
"""
Search the index for relevant chunks using retrieval pipeline
Args:
query: Search query
top_k: Number of results to return
mode: Search mode (uses config default if not specified)
Returns:
List of search results
"""
mode = mode or self.config.mode
# Map IndexMode to retrieval pipeline mode strings
mode_map = {
IndexMode.DENSE: "dense",
IndexMode.SPARSE: "sparse",
IndexMode.HYBRID: "hybrid"
}
search_mode = mode_map.get(mode, "hybrid")
if not top_k or top_k < 1:
top_k = 3
try:
# Query the retrieval pipeline
# Note: The pipeline has two top_k parameters:
# - top_k: for initial retrieval (we set to max(20, top_k))
# - rerank_top_k: for final results (we set to the requested top_k)
response = requests.post(
f"{self.retrieval_url}/search",
json={
"query": query,
"mode": search_mode,
"top_k": max(20, top_k), # Retrieve more candidates for better reranking
"rerank_top_k": top_k, # Return the requested number of results
"skip_reranking": False # Always use reranking for better quality
}
)
response.raise_for_status()
data = response.json()
# Get results based on mode
if search_mode == "hybrid" and "reranked_results" in data:
search_results = data["reranked_results"]
elif search_mode == "dense" and "dense_results" in data:
search_results = data["dense_results"]
elif search_mode == "sparse" and "sparse_results" in data:
search_results = data["sparse_results"]
else:
# Fallback to any available results
search_results = (data.get("reranked_results", []) or
data.get("dense_results", []) or
data.get("sparse_results", []))
# Convert to SearchResult objects
results = []
for item in search_results:
# Try to get our chunk_id from different sources
chunk_id = None
# First, check if metadata contains our doc_id
metadata = item.get("metadata", {})
if metadata.get("doc_id"):
chunk_id = metadata.get("doc_id")
else:
# Fall back to doc_id mapping
generated_doc_id = item.get("doc_id", "")
chunk_id = self.doc_id_mapping.get(generated_doc_id)
# Get chunk from local storage
if chunk_id and chunk_id in self.chunks:
chunk = self.chunks[chunk_id]
# Get score based on result type
score = item.get("rerank_score", item.get("score", 0.0))
results.append(SearchResult(
chunk_id=chunk_id,
score=float(score),
chunk=chunk,
match_type=search_mode
))
else:
# Log warning but don't fail
doc_id = item.get("doc_id", "unknown")
if chunk_id:
logger.debug(f"Chunk {chunk_id} not found in local storage")
else:
logger.debug(f"No mapping found for doc_id {doc_id}")
logger.info(f"Search returned {len(results)} results from retrieval pipeline")
return results
except requests.exceptions.RequestException as e:
logger.error(f"Error searching via retrieval pipeline: {e}")
logger.info("Falling back to empty results. Ensure retrieval pipeline is running.")
return []
def save_index(self, path: Optional[str] = None):
"""
Save the chunks and metadata to disk
Args:
path: Path to save index (uses config default if not specified)
"""
path = path or self.config.index_path
# Save chunks
chunks_data = {
chunk_id: chunk.to_dict()
for chunk_id, chunk in self.chunks.items()
}
with open(f"{path}_chunks.json", 'w', encoding='utf-8') as f:
json.dump(chunks_data, f, ensure_ascii=False, indent=2)
# Save chunk texts
with open(f"{path}_texts.json", 'w', encoding='utf-8') as f:
json.dump(self.chunk_texts, f, ensure_ascii=False, indent=2)
logger.info(f"Chunks saved to {path}. Total chunks: {len(self.chunks)}")
def load_index(self, path: Optional[str] = None):
"""
Load chunks from disk and re-index in retrieval pipeline
Args:
path: Path to load index from (uses config default if not specified)
"""
path = path or self.config.index_path
try:
# Load chunks
with open(f"{path}_chunks.json", 'r', encoding='utf-8') as f:
chunks_data = json.load(f)
self.chunks = {}
for chunk_id, chunk_dict in chunks_data.items():
# Convert messages
from chunker import ConversationMessage
messages = []
for msg_data in chunk_dict.get('messages', []):
messages.append(ConversationMessage(**msg_data))
# Create chunk
chunk = ConversationChunk(
chunk_id=chunk_dict['chunk_id'],
conversation_id=chunk_dict['conversation_id'],
test_id=chunk_dict['test_id'],
chunk_index=chunk_dict['chunk_index'],
start_round=chunk_dict['start_round'],
end_round=chunk_dict['end_round'],
messages=messages,
metadata=chunk_dict.get('metadata', {}),
context_before=chunk_dict.get('context_before'),
context_after=chunk_dict.get('context_after'),
created_at=chunk_dict.get('created_at', '')
)
self.chunks[chunk_id] = chunk
# Load chunk texts if available
texts_path = f"{path}_texts.json"
if Path(texts_path).exists():
with open(texts_path, 'r', encoding='utf-8') as f:
self.chunk_texts = json.load(f)
else:
# Regenerate texts if not saved
self.chunk_texts = {}
for chunk_id, chunk in self.chunks.items():
self.chunk_texts[chunk_id] = self._prepare_chunk_text(chunk)
logger.info(f"Loaded {len(self.chunks)} chunks from {path}")
# Re-index in retrieval pipeline
self.build_indexes()
except Exception as e:
logger.error(f"Error loading index: {e}")
raise
main.py¶
#!/usr/bin/env python3
"""Main entry point for Contextual Retrieval + Advanced Memory Cards System
This demonstrates the dual-layer memory system combining:
1. Contextual chunking for conversation history
2. Advanced JSON cards for structured facts
"""
import argparse
import json
import logging
import sys
from pathlib import Path
from typing import Optional, List
from datetime import datetime
from rich.console import Console
from rich.prompt import Prompt, Confirm
from rich.table import Table
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
from config import Config
from contextual_evaluator import ContextualMemoryEvaluator
from contextual_indexer import ContextualMemoryIndexer
from contextual_agent import ContextualUserMemoryAgent
from advanced_memory_manager import AdvancedMemoryCard, create_sample_cards
from chunker import ConversationChunker
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
# Rich console for better output
console = Console()
class InteractiveContextualRAG:
"""Interactive interface for the contextual RAG system"""
def __init__(self, config: Optional[Config] = None):
"""Initialize the interactive system"""
self.config = config or Config.from_env()
self.evaluator = ContextualMemoryEvaluator(self.config)
self.current_user = "demo_user"
self.indexer = None
self.agent = None
def run(self):
"""Run the interactive session"""
console.print(Panel.fit(
"[bold cyan]Contextual RAG + Advanced Memory Cards System[/bold cyan]\n"
"双层记忆系统：上下文感知检索 + 结构化记忆卡片\n"
"[dim]LLM Judge enabled for automatic evaluation[/dim]",
border_style="cyan"
))
while True:
self.show_menu()
choice = Prompt.ask(
"Select an option",
choices=["1", "2", "3", "4", "5", "6", "7", "8", "0"],
default="1"
)
if choice == "1":
self.demo_mode()
elif choice == "2":
self.load_and_index_conversations()
elif choice == "3":
self.manage_memory_cards()
elif choice == "4":
self.test_query()
elif choice == "5":
self.evaluate_test_cases()
elif choice == "6":
self.evaluate_specific_test_case()
elif choice == "7":
self.show_statistics()
elif choice == "8":
self.configure_settings()
elif choice == "0":
if Confirm.ask("Are you sure you want to exit?"):
console.print("[yellow]Goodbye![/yellow]")
break
def show_menu(self):
"""Display the main menu"""
console.print("\n[bold]Main Menu:[/bold]")
console.print("1. 🚀 Demo Mode (Quick Start)")
console.print("2. 📚 Load & Index Conversations")
console.print("3. 🎴 Manage Memory Cards")
console.print("4. 🔍 Test Query")
console.print("5. 📊 Evaluate All Test Cases (by Category) [LLM Judge]")
console.print("6. 🎯 Evaluate Specific Test Case [LLM Judge]")
console.print("7. 📈 Show Statistics")
console.print("8. ⚙️ Configure Settings")
console.print("0. Exit")
def demo_mode(self):
"""Run a quick demo with sample data"""
console.print("\n[cyan]Demo Mode - Quick Start[/cyan]")
# Initialize components
user_id = "demo_user"
self.indexer = ContextualMemoryIndexer(
user_id=user_id,
use_contextual=True
)
# Create sample memory cards
console.print("\n[yellow]Creating sample memory cards...[/yellow]")
sample_cards = create_sample_cards()
for card in sample_cards:
self.indexer.memory_manager.add_card(card)
console.print(f"[green]✓ Added {len(sample_cards)} memory cards[/green]")
# Create sample conversation chunks
console.print("\n[yellow]Creating sample conversation chunks...[/yellow]")
sample_chunks = self._create_sample_chunks()
# Process with contextual chunking
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
console=console
) as progress:
task = progress.add_task("Processing conversations...", total=None)
result = self.indexer.process_conversation_history(
chunks=sample_chunks,
conversation_id="demo_conv",
generate_summary_cards=False
)
progress.update(task, completed=True)
console.print(f"[green]✓ Indexed {result['contextual_chunks']} contextual chunks[/green]")
# Initialize agent
self.agent = ContextualUserMemoryAgent(
indexer=self.indexer,
config=self.config
)
# Show memory status
console.print("\n[bold]Memory System Status:[/bold]")
console.print(f" Memory Cards: {sum(len(cards) for cards in self.indexer.memory_manager.categories.values())}")
console.print(f" Contextual Chunks: {len(self.indexer.contextual_chunks)}")
# Test queries
test_queries = [
"我的护照什么时候过期？",
"我一月份的东京之行需要准备什么？",
"我的医疗信息有哪些？"
]
console.print("\n[bold]Test Queries:[/bold]")
for i, query in enumerate(test_queries, 1):
console.print(f"\n[cyan]Query {i}:[/cyan] {query}")
if Confirm.ask("Run this query?", default=True):
trajectory = self.agent.answer_question(
question=query,
test_id=f"demo_{i}",
stream=False
)
console.print(Panel(
trajectory.final_answer or "No answer generated",
title="Answer",
border_style="green"
))
if trajectory.memory_cards_used:
console.print(f" Memory cards used: {', '.join(trajectory.memory_cards_used)}")
if trajectory.chunks_retrieved:
console.print(f" Chunks retrieved: {len(trajectory.chunks_retrieved)}")
def _create_sample_chunks(self):
"""Create sample conversation chunks for demo"""
from chunker import ConversationChunk, ConversationMessage
chunks = []
# Sample conversation about travel
messages = [
ConversationMessage("user", "我想订一张去东京的机票", 1),
ConversationMessage("assistant", "好的，请问您什么时候出发？", 2),
ConversationMessage("user", "1月25日出发，2月1日返回", 3),
ConversationMessage("assistant", "让我为您查询1月25日到2月1日的东京往返机票", 4),
]
chunk = ConversationChunk(
chunk_id="demo_chunk_001",
conversation_id="demo_conv",
test_id="demo",
chunk_index=0,
start_round=1,
end_round=2,
messages=messages,
metadata={"topic": "travel"}
)
chunks.append(chunk)
# Sample conversation about passport
messages2 = [
ConversationMessage("user", "我的护照快过期了，什么时候需要续签？", 5),
ConversationMessage("assistant", "您的护照将于2025年2月18日过期，建议提前3-6个月办理续签", 6),
ConversationMessage("user", "好的，我会尽快去办理", 7),
ConversationMessage("assistant", "建议您在出国前确保护照有效期至少6个月", 8),
]
chunk2 = ConversationChunk(
chunk_id="demo_chunk_002",
conversation_id="demo_conv",
test_id="demo",
chunk_index=1,
start_round=3,
end_round=4,
messages=messages2,
metadata={"topic": "passport"}
)
chunks.append(chunk2)
return chunks
def load_and_index_conversations(self):
"""Load and index conversation histories"""
console.print("\n[cyan]Load & Index Conversations[/cyan]")
# Get user ID
user_id = Prompt.ask("Enter user ID", default=self.current_user)
self.current_user = user_id
# Initialize indexer
self.indexer = ContextualMemoryIndexer(
user_id=user_id,
use_contextual=Confirm.ask("Enable contextual chunking?", default=True)
)
# Load conversation files
conv_dir = Prompt.ask(
"Enter conversation directory path",
default="../../week2/user-memory-evaluation/conversations"
)
conv_path = Path(conv_dir)
if not conv_path.exists():
console.print(f"[red]Directory not found: {conv_path}[/red]")
return
# Process conversation files
json_files = list(conv_path.glob("*.json"))
console.print(f"Found {len(json_files)} conversation files")
if not json_files:
console.print("[yellow]No JSON files found[/yellow]")
return
# Process each file
chunker = ConversationChunker(self.config.chunking)
all_chunks = []
with Progress(console=console) as progress:
task = progress.add_task("Processing files...", total=len(json_files))
for json_file in json_files:
try:
with open(json_file, 'r', encoding='utf-8') as f:
data = json.load(f)
# Extract conversations
conversations = data if isinstance(data, dict) else {"conv": data}
for conv_id, messages in conversations.items():
chunks = chunker.chunk_conversation(
messages=messages,
conversation_id=conv_id,
test_id=json_file.stem
)
all_chunks.extend(chunks)
progress.advance(task)
except Exception as e:
console.print(f"[red]Error processing {json_file}: {e}[/red]")
console.print(f"[green]Created {len(all_chunks)} chunks[/green]")
# Index with contextual processing
if all_chunks:
result = self.indexer.process_conversation_history(
chunks=all_chunks,
conversation_id="batch_index",
generate_summary_cards=Confirm.ask("Generate summary cards?", default=True)
)
console.print(f"[green]✓ Indexed {result['contextual_chunks']} contextual chunks[/green]")
console.print(f"[green]✓ Total memory cards: {result['memory_cards_after']}[/green]")
def manage_memory_cards(self):
"""Manage advanced memory cards"""
if not self.indexer:
console.print("[yellow]Please initialize the system first (option 1 or 2)[/yellow]")
return
console.print("\n[cyan]Memory Card Management[/cyan]")
# Show current cards
stats = self.indexer.memory_manager.get_statistics()
console.print(f"\nCurrent cards: {stats['total_cards']}")
for category, info in stats['categories'].items():
console.print(f" {category}: {info['count']} cards")
# Options
console.print("\n1. View all cards")
console.print("2. Add new card")
console.print("3. Search cards")
console.print("4. Delete card")
console.print("5. Back")
choice = Prompt.ask("Select option", choices=["1", "2", "3", "4", "5"])
if choice == "1":
# View all cards
context = self.indexer.memory_manager.get_context_string()
console.print(Panel(context, title="Memory Cards", border_style="cyan"))
elif choice == "2":
# Add new card
category = Prompt.ask("Category")
card_key = Prompt.ask("Card key")
backstory = Prompt.ask("Backstory")
person = Prompt.ask("Person", default="User")
relationship = Prompt.ask("Relationship", default="primary")
# Get additional data fields
data = {}
while True:
field = Prompt.ask("Add data field (empty to finish)")
if not field:
break
value = Prompt.ask(f"Value for {field}")
data[field] = value
# Create and add card
card = AdvancedMemoryCard(
category=category,
card_key=card_key,
backstory=backstory,
date_created=datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
person=person,
relationship=relationship,
data=data
)
memory_id = self.indexer.memory_manager.add_card(card)
console.print(f"[green]✓ Added card: {memory_id}[/green]")
elif choice == "3":
# Search cards
query = Prompt.ask("Search query")
results = self.indexer.memory_manager.search_cards(query)
if results:
console.print(f"\n[green]Found {len(results)} cards:[/green]")
for memory_id, card in results:
console.print(f"\n{memory_id}:")
console.print(f" Backstory: {card.backstory}")
console.print(f" Person: {card.person}")
else:
console.print("[yellow]No cards found[/yellow]")
elif choice == "4":
# Delete card
category = Prompt.ask("Category")
card_key = Prompt.ask("Card key")
if Confirm.ask(f"Delete {category}.{card_key}?"):
if self.indexer.memory_manager.delete_card(category, card_key):
console.print("[green]✓ Card deleted[/green]")
else:
console.print("[red]Card not found[/red]")
def test_query(self):
"""Test a query against the system"""
if not self.indexer:
console.print("[yellow]Please initialize the system first (option 1 or 2)[/yellow]")
return
if not self.agent:
self.agent = ContextualUserMemoryAgent(
indexer=self.indexer,
config=self.config
)
console.print("\n[cyan]Test Query[/cyan]")
# Show current memory status
console.print(f"\nMemory Status:")
console.print(f" Cards: {sum(len(cards) for cards in self.indexer.memory_manager.categories.values())}")
console.print(f" Chunks: {len(self.indexer.contextual_chunks)}")
# Get query
query = Prompt.ask("\nEnter your question")
# Process query
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
console=console
) as progress:
task = progress.add_task("Processing...", total=None)
trajectory = self.agent.answer_question(
question=query,
test_id="interactive",
stream=False
)
progress.update(task, completed=True)
# Display results
console.print(Panel(
trajectory.final_answer or "No answer generated",
title="Answer",
border_style="green"
))
# Show details
console.print(f"\n[bold]Query Details:[/bold]")
console.print(f" Iterations: {len(trajectory.iterations)}")
console.print(f" Tool calls: {len(trajectory.tool_calls)}")
if trajectory.memory_cards_used:
console.print(f"\n[bold]Memory Cards Used:[/bold]")
for card_id in trajectory.memory_cards_used:
console.print(f" • {card_id}")
if trajectory.chunks_retrieved:
console.print(f"\n[bold]Chunks Retrieved:[/bold] {len(trajectory.chunks_retrieved)}")
if Confirm.ask("Show chunk details?"):
for chunk_id in trajectory.chunks_retrieved[:3]:
if chunk_id in self.indexer.contextual_chunks:
chunk = self.indexer.contextual_chunks[chunk_id]
console.print(f"\n Chunk: {chunk_id}")
console.print(f" Context: {chunk.context[:200]}...")
def evaluate_specific_test_case(self):
"""Evaluate a specific test case selected by the user"""
console.print("\n[cyan]Evaluate Specific Test Case[/cyan]")
# First, load all test cases to show to the user
console.print("\nLoading available test cases...")
# Load all categories
all_test_cases = []
categories = ["layer1", "layer2", "layer3"]
for category in categories:
test_cases = self.evaluator.load_test_cases(category)
for test_id in test_cases:
test_case = self.evaluator.test_cases[test_id]
all_test_cases.append({
"id": test_id,
"category": category,
"title": test_case.title,
"conversations": len(test_case.conversation_histories)
})
if not all_test_cases:
console.print("[yellow]No test cases found[/yellow]")
return
# Sort test cases by test ID (name)
all_test_cases.sort(key=lambda x: x["id"])
console.print(f"\n[green]Found {len(all_test_cases)} test cases[/green]")
# Create a table to display test cases
table = Table(title="Available Test Cases (Sorted by Name)", show_lines=True)
table.add_column("#", style="dim", width=4)
table.add_column("Test ID", style="cyan", width=25)
table.add_column("Category", style="magenta", width=8)
table.add_column("Title", style="green", width=50)
table.add_column("Conv.", justify="right", width=5)
for idx, test_info in enumerate(all_test_cases, 1):
title = test_info["title"][:47] + "..." if len(test_info["title"]) > 50 else test_info["title"]
table.add_row(
str(idx),
test_info["id"],
test_info["category"],
title,
str(test_info["conversations"])
)
console.print(table)
# Let user select a test case
console.print("\n[bold]Select a test case to evaluate:[/bold]")
console.print("Enter the number (#) or the Test ID directly")
user_input = Prompt.ask("Your choice")
# Find the selected test case
selected_test_id = None
# Check if user entered a number
if user_input.isdigit():
idx = int(user_input) - 1
if 0 <= idx < len(all_test_cases):
selected_test_id = all_test_cases[idx]["id"]
else:
console.print(f"[red]Invalid number: {user_input}[/red]")
return
else:
# Check if user entered a test ID
for test_info in all_test_cases:
if test_info["id"] == user_input:
selected_test_id = user_input
break
if not selected_test_id:
console.print(f"[red]Test case not found: {user_input}[/red]")
return
# Get the test case details
test_case = self.evaluator.test_cases[selected_test_id]
# Show test case details
console.print(Panel(
f"[bold]{test_case.title}[/bold]\n\n"
f"Category: {test_case.category}\n"
f"Description: {test_case.description}\n\n"
f"[yellow]User Question:[/yellow]\n{test_case.user_question}\n\n"
f"[green]Evaluation Criteria:[/green]\n{test_case.evaluation_criteria[:200]}...\n\n"
f"Conversations: {len(test_case.conversation_histories)}",
title=selected_test_id,
border_style="cyan"
))
# Run evaluation
console.print(f"\n[cyan]Evaluating {selected_test_id}...[/cyan]")
console.print(f"[dim]Using LLM Judge for automatic evaluation[/dim]\n")
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
console=console
) as progress:
task = progress.add_task("Processing...", total=None)
try:
result = self.evaluator.evaluate_test_case(selected_test_id)
progress.update(task, completed=True)
# Display result
status = "✓ Success" if result.success else "✗ Failed"
console.print(f"\n[{'green' if result.success else 'red'}]{status}[/{'green' if result.success else 'red'}]")
console.print("\n[bold]Agent Answer:[/bold]")
console.print(Panel(result.agent_answer or "No answer generated", border_style="cyan"))
console.print("\n[bold]Evaluation Criteria:[/bold]")
console.print(Panel(result.evaluation_criteria, border_style="green"))
# Display LLM evaluation if available
if result.llm_evaluation:
console.print("\n[bold cyan]LLM Judge Evaluation:[/bold cyan]")
llm_eval = result.llm_evaluation
reward = llm_eval.get('reward', 0)
passed = llm_eval.get('passed', False)
# Format reward with color based on score
if reward >= 0.8:
reward_color = "green"
elif reward >= 0.6:
reward_color = "yellow"
else:
reward_color = "red"
console.print(f" Reward Score: [{reward_color}]{reward:.3f}/1.000[/{reward_color}]")
console.print(f" Passed: [{'green' if passed else 'red'}]{'Yes' if passed else 'No'}[/{'green' if passed else 'red'}]")
if 'reasoning' in llm_eval:
console.print(f"\n[bold]Reasoning:[/bold]")
console.print(Panel(llm_eval['reasoning'], border_style="cyan"))
if 'required_info_found' in llm_eval and llm_eval['required_info_found']:
console.print(f"\n[bold]Required Information Found:[/bold]")
for key, found in llm_eval['required_info_found'].items():
status = "✓" if found else "✗"
color = "green" if found else "red"
console.print(f" [{color}]{status}[/{color}] {key}")
console.print(f"\n[bold]Statistics:[/bold]")
console.print(f" Iterations: {result.iterations}")
console.print(f" Tool Calls: {result.tool_calls}")
console.print(f" Memory Cards Used: {len(result.memory_cards_used)}")
console.print(f" Chunks Retrieved: {len(result.chunks_retrieved)}")
console.print(f" Contextual Chunks: {result.contextual_chunks_count}")
console.print(f" Processing Time: {result.processing_time:.2f}s")
console.print(f" Context Generation Time: {result.context_generation_time:.2f}s")
if result.error:
console.print(f"\n[red]Error: {result.error}[/red]")
except Exception as e:
progress.update(task, completed=True)
console.print(f"[red]Error evaluating test case: {e}[/red]")
def evaluate_test_cases(self):
"""Run evaluation on all test cases in a category"""
console.print("\n[cyan]Evaluate All Test Cases (by Category)[/cyan]")
# Load test cases
category = Prompt.ask(
"Select category",
choices=["all", "layer1", "layer2", "layer3"],
default="layer1"
)
test_cases = self.evaluator.load_test_cases(
category=None if category == "all" else category
)
# Sort test cases by ID
test_cases = sorted(test_cases)
console.print(f"[green]Loaded {len(test_cases)} test cases (sorted by name)[/green]")
console.print(f"[dim]Using LLM Judge for automatic evaluation[/dim]\n")
# Run evaluation
if Confirm.ask("Run evaluation?"):
with Progress(console=console) as progress:
task = progress.add_task("Evaluating...", total=len(test_cases))
for test_id in test_cases:
try:
result = self.evaluator.evaluate_test_case(test_id)
progress.advance(task)
except Exception as e:
console.print(f"[red]Error evaluating {test_id}: {e}[/red]")
progress.advance(task)
# Show results
report = self.evaluator.generate_report()
console.print("\n" + report)
# Save results
if Confirm.ask("Save results to file?"):
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = f"results/evaluation_{timestamp}.json"
self.evaluator.save_results(output_file)
console.print(f"[green]✓ Results saved to {output_file}[/green]")
def show_statistics(self):
"""Show system statistics"""
console.print("\n[cyan]System Statistics[/cyan]")
if self.indexer:
stats = self.indexer.get_statistics()
# Create statistics table
table = Table(title="Contextual Memory Statistics")
table.add_column("Metric", style="cyan")
table.add_column("Value", justify="right")
# Indexer stats
table.add_row("Indexed Chunks", str(stats.get("chunks_indexed", 0)))
table.add_row("Memory Cards", str(stats.get("memory_cards", 0)))
table.add_row("Indexing Time", f"{stats.get('indexing_time', 0):.2f}s")
# Chunker stats
if "chunker_stats" in stats:
cs = stats["chunker_stats"]
table.add_row("Contextual Chunks", str(cs.get("contextual_chunks", 0)))
table.add_row("Context Tokens", str(cs.get("total_context_tokens", 0)))
table.add_row("Cache Hit Rate", f"{cs.get('cache_hit_rate', 0):.1%}")
table.add_row("Est. Cost", f"${cs.get('estimated_cost', 0):.3f}")
# Memory stats
if "memory_stats" in stats:
ms = stats["memory_stats"]
table.add_row("Total Cards", str(ms.get("total_cards", 0)))
for cat, info in ms.get("categories", {}).items():
table.add_row(f" {cat}", str(info.get("count", 0)))
console.print(table)
else:
console.print("[yellow]System not initialized[/yellow]")
def configure_settings(self):
"""Configure system settings"""
console.print("\n[cyan]Configuration Settings[/cyan]")
# Show current settings
console.print(f"\nCurrent Settings:")
console.print(f" LLM Provider: {self.config.llm.provider}")
console.print(f" LLM Model: {self.config.llm.model}")
console.print(f" Chunking: {self.config.chunking.rounds_per_chunk} rounds/chunk")
console.print(f" Index Mode: {self.config.index.mode}")
if Confirm.ask("\nModify settings?"):
# LLM settings
if Confirm.ask("Change LLM provider?"):
provider = Prompt.ask(
"Provider",
choices=["kimi", "doubao", "siliconflow", "openai"],
default=self.config.llm.provider
)
self.config.llm.provider = provider
# Chunking settings
if Confirm.ask("Change chunking settings?"):
rounds = Prompt.ask(
"Rounds per chunk",
default=str(self.config.chunking.rounds_per_chunk)
)
self.config.chunking.rounds_per_chunk = int(rounds)
console.print("[green]✓ Settings updated[/green]")
def main():
"""主入口：实验 3-12 上下文感知检索增强用户记忆"""
parser = argparse.ArgumentParser(
description=(
"实验 3-12：利用上下文感知检索增强用户记忆。\n"
"在把对话记忆块送入嵌入/索引前先生成『上下文前缀』，"
"提升脱离上下文的孤立片段（如『好的，就订这个吧』）的召回。"
),
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=(
"示例：\n"
" python main.py --mode compare # 离线对比上下文化 vs 原始块（无需 API）\n"
" python main.py --mode compare --query '我的护照什么时候过期？' # 单条查询离线检索对比\n"
" python main.py --mode compare --output results/compare.json # 保存对比结果\n"
" python main.py --mode evaluate --category layer1 # 端到端评估（需 API/检索服务）\n"
" python main.py --mode interactive # 交互式界面（默认，需 API）\n"
),
)
parser.add_argument(
"--mode",
choices=["interactive", "evaluate", "demo", "compare"],
default="interactive",
help="运行模式：interactive 交互式(默认) / evaluate 端到端评估 / demo 演示 / compare 离线对比(无需 API)",
)
parser.add_argument(
"--category",
choices=["layer1", "layer2", "layer3"],
help="评估的测试分类（layer1 基础回忆 / layer2 多会话检索 / layer3 主动服务）",
)
parser.add_argument(
"--config",
type=str,
help="配置文件（JSON）路径",
)
# 离线对比（compare 模式）相关参数
parser.add_argument(
"--dataset",
type=str,
default=None,
help="compare 模式使用的记忆问答对照集 JSON（默认：memory_qa_eval.json）",
)
parser.add_argument(
"--query",
type=str,
default=None,
help="compare 模式下对单条查询做离线检索对比（plain vs contextual 的 Top-K）",
)
parser.add_argument(
"--output",
type=str,
default=None,
help="将 compare / evaluate 的结果保存为 JSON 的路径",
)
# 配置覆盖项（可选，覆盖环境变量/配置文件；不改变默认行为）
parser.add_argument(
"--user-id",
type=str,
default=None,
help="用户标识（写入输出结果作为标签，便于区分多用户记忆）",
)
parser.add_argument(
"--model",
type=str,
default=None,
help="覆盖 LLM 模型名（默认取环境变量/提供商默认值）",
)
parser.add_argument(
"--provider",
type=str,
default=None,
help="覆盖 LLM 提供商（kimi / doubao / siliconflow / openai 等）",
)
parser.add_argument(
"--store-path",
type=str,
default=None,
help="记忆块存储（chunk_store）路径，覆盖默认 data/chunk_store.json",
)
contextual_group = parser.add_mutually_exclusive_group()
contextual_group.add_argument(
"--contextual",
dest="contextual",
action="store_true",
default=None,
help="启用上下文化（索引前为每块生成上下文前缀，默认开启）",
)
contextual_group.add_argument(
"--no-contextual",
dest="contextual",
action="store_false",
help="关闭上下文化（直接索引原始对话块，用于对照）",
)
args = parser.parse_args()
# compare 模式：完全离线，无需加载 LLM / 检索服务配置
if args.mode == "compare":
from contextual_compare import (
run_comparison,
single_query,
DEFAULT_DATASET,
)
dataset = args.dataset or DEFAULT_DATASET
if args.query:
single_query(dataset, args.query)
else:
run_comparison(dataset, output_path=args.output)
return
# Load configuration
if args.config:
config = Config.load(args.config)
else:
config = Config.from_env()
# 应用命令行覆盖项
if args.provider:
config.llm.provider = args.provider
if args.model:
config.llm.model = args.model
if args.store_path:
config.index.chunk_store_path = args.store_path
if args.contextual is not None:
config.index.enable_contextual = args.contextual
if args.mode == "interactive":
# Interactive mode
app = InteractiveContextualRAG(config)
app.run()
elif args.mode == "evaluate":
# Evaluation mode
evaluator = ContextualMemoryEvaluator(config)
test_cases = evaluator.load_test_cases(args.category)
console.print(f"[cyan]Evaluating {len(test_cases)} test cases[/cyan]")
for test_id in test_cases:
try:
result = evaluator.evaluate_test_case(test_id)
status = "✓" if result.success else "✗"
console.print(f"{status} {test_id}: {result.processing_time:.2f}s")
except Exception as e:
console.print(f"✗ {test_id}: Error - {e}")
# Generate report
report = evaluator.generate_report()
console.print("\n" + report)
# Save results
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_file = f"results/evaluation_{timestamp}.json"
Path("results").mkdir(exist_ok=True)
evaluator.save_results(output_file)
console.print(f"[green]Results saved to {output_file}[/green]")
elif args.mode == "demo":
# Demo mode
app = InteractiveContextualRAG(config)
app.demo_mode()
if __name__ == "__main__":
main()
quickstart.py¶
#!/usr/bin/env python3
"""Quick start script for Agentic RAG User Memory Evaluation
This script provides a simple demo to get started with the system.
"""
import os
import sys
from pathlib import Path
from rich.console import Console
from rich.panel import Panel
# Check for required environment variables
console = Console()
def check_environment():
"""Check if required environment variables are set"""
required_vars = []
optional_vars = []
# Check for OpenAI API key (required for embeddings)
if not os.getenv("OPENAI_API_KEY"):
required_vars.append("OPENAI_API_KEY (required for embeddings)")
# Check for at least one LLM provider
llm_providers = [
"KIMI_API_KEY",
"SILICONFLOW_API_KEY",
"DOUBAO_API_KEY",
"OPENROUTER_API_KEY"
]
if not any(os.getenv(key) for key in llm_providers):
required_vars.append("At least one LLM provider API key (KIMI_API_KEY recommended)")
if required_vars:
console.print(Panel(
"[bold red]Missing Required Environment Variables[/bold red]\n\n" +
"\n".join(f"• {var}" for var in required_vars) +
"\n\n[yellow]Please set up your .env file:[/yellow]\n" +
"1. Copy env.example to .env\n" +
"2. Add your API keys\n" +
"3. Run this script again",
border_style="red"
))
return False
return True
def run_quick_demo():
"""Run a quick demonstration"""
from config import Config
from evaluator import UserMemoryEvaluator
console.print(Panel.fit(
"[bold cyan]Agentic RAG for User Memory - Quick Start Demo[/bold cyan]\n"
"This demo will:\n"
"1. Load a simple test case\n"
"2. Chunk the conversation history\n"
"3. Build a RAG index\n"
"4. Answer a question using the indexed memory",
border_style="cyan"
))
console.print("\n[yellow]Initializing system...[/yellow]")
# Create configuration with demo settings
config = Config.from_env()
config.chunking.rounds_per_chunk = 10 # Smaller chunks for demo
config.evaluation.max_iterations = 5 # Fewer iterations for speed
config.agent.enable_reasoning = True # Show reasoning process
# Initialize evaluator
evaluator = UserMemoryEvaluator(config)
# Load test cases (just layer1 for demo)
console.print("\n[yellow]Loading test cases...[/yellow]")
test_cases = evaluator.load_test_cases("layer1")
if not test_cases:
console.print("[red]No test cases found. Please check the path to week2/user-memory-evaluation[/red]")
return
# Use the first test case
test_case = test_cases[0]
test_id = test_case.test_id
console.print(f"\n[green]Selected test case:[/green] {test_case.title}")
console.print(f"[green]Question:[/green] {test_case.user_question}\n")
# Evaluate the test case
console.print("[yellow]Processing conversation history...[/yellow]")
console.print("• Chunking conversations into segments")
console.print("• Building search indexes")
console.print("• Preparing RAG agent\n")
result = evaluator.evaluate_test_case(test_id)
# Display results
console.print("\n" + "="*60)
console.print("[bold green]Demo Results[/bold green]")
console.print("="*60)
console.print(f"\n[bold]Agent's Answer:[/bold]")
console.print(Panel(result.agent_answer, border_style="cyan"))
console.print(f"\n[bold]Expected Answer:[/bold]")
console.print(Panel(result.expected_answer, border_style="green"))
console.print(f"\n[bold]Performance Metrics:[/bold]")
console.print(f"• Success: {'✓ Yes' if result.success else '✗ No'}")
console.print(f"• Iterations: {result.iterations}")
console.print(f"• Tool Calls: {result.tool_calls}")
console.print(f"• Chunks Created: {result.chunk_count}")
console.print(f"• Processing Time: {result.processing_time:.2f} seconds")
console.print(f"• Indexing Time: {result.indexing_time:.2f} seconds")
console.print("\n[bold cyan]Demo Complete![/bold cyan]")
console.print("\nTo explore more:")
console.print("• Run [bold]python main.py[/bold] for interactive mode")
console.print("• Run [bold]python main.py --mode batch --category layer1[/bold] for batch evaluation")
console.print("• Check the README.md for detailed documentation")
def main():
"""Main entry point"""
console.print("\n[bold]Agentic RAG for User Memory Evaluation - Quick Start[/bold]\n")
# Check environment
if not check_environment():
sys.exit(1)
# Check if .env file exists
if not Path(".env").exists() and Path("env.example").exists():
console.print("[yellow]Creating .env file from env.example...[/yellow]")
import shutil
shutil.copy("env.example", ".env")
console.print("[red]Please edit .env file with your API keys and run again.[/red]")
sys.exit(1)
# Load environment variables
from dotenv import load_dotenv
load_dotenv()
# Run the demo
try:
run_quick_demo()
except KeyboardInterrupt:
console.print("\n[yellow]Demo interrupted by user[/yellow]")
except Exception as e:
console.print(f"\n[red]Error during demo: {e}[/red]")
console.print("[yellow]Please check your configuration and try again[/yellow]")
import traceback
if os.getenv("DEBUG"):
traceback.print_exc()
if __name__ == "__main__":
main()
test_contextual_system.py¶
#!/usr/bin/env python3
"""Test script for the Contextual Retrieval + Advanced Memory Cards System"""
import logging
from config import Config
from contextual_evaluator import ContextualMemoryEvaluator
from contextual_indexer import ContextualMemoryIndexer
from contextual_agent import ContextualUserMemoryAgent
from advanced_memory_manager import create_sample_cards
from chunker import ConversationChunk, ConversationMessage
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def test_dual_memory_system():
"""Test the dual memory system with a sample scenario"""
print("\n" + "="*60)
print("Testing Contextual Retrieval + Advanced Memory Cards")
print("="*60)
# Initialize components
config = Config.from_env()
user_id = "test_user_contextual"
# Create indexer with contextual chunking
print("\n1. Initializing Contextual Memory Indexer...")
indexer = ContextualMemoryIndexer(
user_id=user_id,
use_contextual=True
)
print(f" ✓ Indexer initialized")
# Add sample memory cards
print("\n2. Adding Advanced Memory Cards...")
sample_cards = create_sample_cards()
for card in sample_cards:
indexer.memory_manager.add_card(card)
print(f" ✓ Added {len(sample_cards)} memory cards")
# Create sample conversation chunks
print("\n3. Creating Sample Conversation Chunks...")
chunks = []
# Conversation about travel
messages1 = [
ConversationMessage("user", "我想订一张去东京的机票", 1),
ConversationMessage("assistant", "好的，请问您什么时候出发？", 2),
ConversationMessage("user", "1月25日出发，2月1日返回", 3),
ConversationMessage("assistant", "让我为您查询1月25日到2月1日的东京往返机票", 4),
]
chunk1 = ConversationChunk(
chunk_id="test_chunk_001",
conversation_id="test_conv",
test_id="test",
chunk_index=0,
start_round=1,
end_round=2,
messages=messages1,
metadata={"topic": "travel"}
)
chunks.append(chunk1)
# Conversation about passport
messages2 = [
ConversationMessage("user", "我的护照快过期了，什么时候需要续签？", 5),
ConversationMessage("assistant", "您的护照将于2025年2月18日过期，建议提前3-6个月办理续签", 6),
ConversationMessage("user", "好的，我会尽快去办理", 7),
ConversationMessage("assistant", "建议您在出国前确保护照有效期至少6个月", 8),
]
chunk2 = ConversationChunk(
chunk_id="test_chunk_002",
conversation_id="test_conv",
test_id="test",
chunk_index=1,
start_round=3,
end_round=4,
messages=messages2,
metadata={"topic": "passport"}
)
chunks.append(chunk2)
print(f" ✓ Created {len(chunks)} conversation chunks")
# Process with contextual chunking
print("\n4. Processing with Contextual Chunking...")
result = indexer.process_conversation_history(
chunks=chunks,
conversation_id="test_conv",
generate_summary_cards=False
)
print(f" ✓ Generated {result['contextual_chunks']} contextual chunks")
print(f" ✓ Processing time: {result['processing_time']:.2f}s")
# Initialize agent
print("\n5. Initializing Contextual Agent...")
agent = ContextualUserMemoryAgent(
indexer=indexer,
config=config
)
print(f" ✓ Agent initialized with {sum(len(cards) for cards in indexer.memory_manager.categories.values())} memory cards")
# Test queries
print("\n6. Testing Queries...")
test_queries = [
("我的护照什么时候过期？", "Should find passport expiration date from memory cards"),
("我一月份的东京之行需要准备什么？", "Should combine travel and passport info"),
("我的银行账户信息是什么？", "Should find bank account from memory cards"),
]
for i, (query, expected) in enumerate(test_queries, 1):
print(f"\n Query {i}: {query}")
print(f" Expected: {expected}")
trajectory = agent.answer_question(
question=query,
test_id=f"test_{i}",
stream=False
)
if trajectory.final_answer:
print(f" Answer: {trajectory.final_answer[:200]}...")
print(f" ✓ Memory cards used: {len(trajectory.memory_cards_used)}")
print(f" ✓ Chunks retrieved: {len(trajectory.chunks_retrieved)}")
else:
print(f" ✗ No answer generated")
# Show statistics
print("\n7. System Statistics:")
stats = indexer.get_statistics()
print(f" • Chunks indexed: {stats.get('chunks_indexed', 0)}")
print(f" • Memory cards: {stats.get('memory_cards', 0)}")
if 'chunker_stats' in stats:
cs = stats['chunker_stats']
print(f" • Context generation tokens: {cs.get('total_context_tokens', 0)}")
print(f" • Estimated cost: ${cs.get('estimated_cost', 0):.3f}")
print("\n" + "="*60)
print("Test Complete! The dual memory system is working correctly.")
print("="*60)
def test_evaluation_system():
"""Test the evaluation system with Layer 1 test cases"""
print("\n" + "="*60)
print("Testing Evaluation System")
print("="*60)
config = Config.from_env()
evaluator = ContextualMemoryEvaluator(config)
# Load Layer 1 test cases
print("\n1. Loading Test Cases...")
test_cases = evaluator.load_test_cases("layer1")
print(f" ✓ Loaded {len(test_cases)} test cases")
if test_cases:
# Test the first case
first_test = test_cases[0]
print(f"\n2. Testing First Case: {first_test}")
test_case = evaluator.test_cases[first_test]
print(f" Title: {test_case.title}")
print(f" Category: {test_case.category}")
print(f" Conversations: {len(test_case.conversation_histories)}")
# Run evaluation
print("\n3. Running Evaluation...")
try:
result = evaluator.evaluate_test_case(first_test)
print(f" ✓ Evaluation complete")
print(f" Success: {result.success}")
print(f" Iterations: {result.iterations}")
print(f" Tool calls: {result.tool_calls}")
print(f" Processing time: {result.processing_time:.2f}s")
if result.agent_answer:
print(f" Answer preview: {result.agent_answer[:100]}...")
except Exception as e:
print(f" ✗ Evaluation failed: {e}")
print("\n" + "="*60)
print("Evaluation System Test Complete!")
print("="*60)
if __name__ == "__main__":
import sys
if len(sys.argv) > 1 and sys.argv[1] == "eval":
test_evaluation_system()
else:
test_dual_memory_system()
test_enhanced_logging.py¶
#!/usr/bin/env python3
"""Test enhanced logging of LLM responses"""
import os
import sys
import logging
# Set up logging to see all messages
logging.basicConfig(
level=logging.INFO,
format='%(levelname)s:%(name)s:%(message)s'
)
# Ensure we have API keys
if not os.getenv("KIMI_API_KEY"):
print("Please set KIMI_API_KEY environment variable")
sys.exit(1)
from config import Config
from agent import UserMemoryRAGAgent
def test_logging():
"""Test that LLM responses are properly logged"""
print("\n" + "="*80)
print("Testing Enhanced LLM Response Logging")
print("="*80)
# Initialize agent
config = Config.from_env()
agent = UserMemoryRAGAgent(config)
# Test with a simple question
test_question = "What is the purpose of this system?"
print(f"\nTest Question: {test_question}")
print("\nYou should now see:")
print("1. Iteration info")
print("2. LLM Response content (up to 500 chars)")
print("3. Tool calls if any")
print("4. Final answer")
print("\n" + "-"*80)
# Run the agent
result = agent.answer_question(
question=test_question,
test_id="test_logging",
stream=False
)
print("\n" + "-"*80)
print(f"\nFinal Answer: {result.get('answer', 'No answer')[:200]}...")
print(f"Success: {result.get('success', False)}")
print(f"Iterations: {result.get('iterations', 0)}")
print(f"Tool Calls: {result.get('tool_calls', 0)}")
print("\n✓ Enhanced logging is working!")
print(" - LLM responses are shown during iterations")
print(" - Tool calls and results are logged")
print(" - Evaluation reasoning will be shown when running tests")
print("="*80)
if __name__ == "__main__":
test_logging()
test_fallback_search_empty_query.py¶
"""Regression test: the fallback local search in search_with_context must not
raise ZeroDivisionError on an empty query or on chunks whose contextualized_text
is empty (e.g. loaded from disk with a missing field)."""
import os
import sys
import types
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from contextual_chunking import ContextualConversationChunk
from contextual_indexer import ContextualMemoryIndexer
def _make_indexer(chunks):
"""Build an indexer without __init__; force the retrieval-pipeline call to
fail so the local-search fallback runs."""
indexer = ContextualMemoryIndexer.__new__(ContextualMemoryIndexer)
indexer.retrieval_url = "http://127.0.0.1:1" # nothing listening -> fallback
indexer.contextual_chunks = {c.chunk_id: c for c in chunks}
indexer.memory_manager = types.SimpleNamespace(search_cards=lambda q: [])
return indexer
def _chunk(chunk_id, contextualized_text):
return ContextualConversationChunk(
chunk_id=chunk_id,
conversation_id="conv1",
test_id="t1",
chunk_index=0,
start_round=0,
end_round=1,
messages=[],
original_text="",
context="",
contextualized_text=contextualized_text,
)
def test_fallback_search_empty_query_and_empty_chunk_text():
indexer = _make_indexer([_chunk("c1", "")])
results = indexer.search_with_context("", top_k=3) # must not raise
assert results["chunk_results"] == []
def test_fallback_search_normal_query_still_matches():
indexer = _make_indexer([_chunk("c1", "the user likes blue shoes")])
results = indexer.search_with_context("blue", top_k=3)
assert len(results["chunk_results"]) == 1
assert results["chunk_results"][0]["chunk_id"] == "c1"
test_fixes.py¶
#!/usr/bin/env python3
"""Test script to verify all fixes are working"""
import logging
from config import Config
from contextual_evaluator import ContextualMemoryEvaluator
# Set up detailed logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
def test_single_evaluation():
"""Test evaluation with a single simple test case"""
print("="*80)
print("TESTING CONTEXTUAL RETRIEVAL SYSTEM - VERBOSE MODE")
print("="*80)
config = Config.from_env()
evaluator = ContextualMemoryEvaluator(config)
# Load only layer1 test cases
test_cases = evaluator.load_test_cases("layer1")
print(f"\nLoaded {len(test_cases)} test cases")
if test_cases:
# Pick the first test case
test_id = test_cases[0]
test_case = evaluator.test_cases[test_id]
print(f"\n{'='*80}")
print(f"EVALUATING: {test_id}")
print(f"Title: {test_case.title}")
print(f"Question: {test_case.user_question}")
print(f"Conversations: {len(test_case.conversation_histories)}")
if test_case.conversation_histories:
first_conv = test_case.conversation_histories[0]
print(f"First conversation has {len(first_conv.get('messages', []))} messages")
print(f"{'='*80}\n")
# Run evaluation
try:
result = evaluator.evaluate_test_case(test_id)
print(f"\n{'='*80}")
print("EVALUATION RESULT")
print(f"{'='*80}")
print(f"Success: {result.success}")
print(f"Iterations: {result.iterations}")
print(f"Tool Calls: {result.tool_calls}")
print(f"Memory Cards Used: {len(result.memory_cards_used)}")
print(f"Chunks Retrieved: {len(result.chunks_retrieved)}")
print(f"Processing Time: {result.processing_time:.2f}s")
if result.agent_answer:
print(f"\nAgent Answer:")
print(result.agent_answer)
if result.error:
print(f"\nError: {result.error}")
print(f"{'='*80}\n")
except Exception as e:
print(f"\nERROR during evaluation: {e}")
import traceback
traceback.print_exc()
else:
print("No test cases available")
if __name__ == "__main__":
test_single_evaluation()
test_llm_evaluation.py¶
#!/usr/bin/env python3
"""Test script to demonstrate LLM evaluation integration"""
import os
import sys
import logging
from pathlib import Path
from rich.console import Console
from rich.panel import Panel
# Set up logging to see evaluation logs
logging.basicConfig(
level=logging.INFO,
format='%(levelname)s:%(name)s:%(message)s'
)
# Set dummy API key for demo
os.environ["KIMI_API_KEY"] = "test_key"
console = Console()
def test_llm_evaluation_integration():
"""Test that LLM evaluation is properly integrated"""
console.print("\n[bold cyan]Testing LLM Evaluation Integration[/bold cyan]")
console.print("="*80)
from config import Config
from evaluator import UserMemoryEvaluator, TestCase, EvaluationResult
# Initialize evaluator
config = Config.from_env()
evaluator = UserMemoryEvaluator(config)
# Check if LLM evaluator was initialized
if evaluator.llm_evaluator:
console.print("[green]✓ LLM Evaluator successfully initialized[/green]")
console.print(" The system will automatically evaluate agent responses")
else:
console.print("[yellow]⚠ LLM Evaluator not available[/yellow]")
console.print(" Automatic evaluation will be skipped")
console.print(" To enable, ensure week2/user-memory-evaluation is accessible")
console.print(" and has proper API keys configured")
# Create a test case for demonstration
test_case = TestCase(
test_id="demo_test",
category="demo",
title="Demo Test Case",
description="Test case for demonstrating LLM evaluation",
conversation_histories=[{
"conversation_id": "demo_conv",
"messages": [
{"role": "user", "content": "What's my account number?"},
{"role": "assistant", "content": "Your account number is 123456789."}
],
"metadata": {"business": "Demo Bank"}
}],
user_question="What is my account number?",
expected_answer="Your account number is 123456789.",
required_information=["account number: 123456789"]
)
# Simulate an agent response
agent_answer = "Based on the conversation history, your account number is 123456789."
console.print("\n[bold]Simulating Evaluation Process:[/bold]")
console.print(f"Test Question: {test_case.user_question}")
console.print(f"Agent Answer: {agent_answer}")
console.print(f"Expected Answer: {test_case.expected_answer}")
if evaluator.llm_evaluator:
console.print("\n[yellow]LLM Evaluation would be triggered automatically when running a test case[/yellow]")
console.print("The evaluation will:")
console.print(" 1. Send the agent's response to the LLM evaluator")
console.print(" 2. Get a reward score (0.0 to 1.0)")
console.print(" 3. Determine if the response passed (reward >= 0.6)")
console.print(" 4. Provide reasoning for the evaluation")
console.print(" 5. Check if required information was found")
console.print(" 6. Display all results in the console")
console.print("\n[bold]Evaluation Flow:[/bold]")
console.print("1. Agent generates response using RAG")
console.print("2. Response is automatically evaluated by LLM")
console.print("3. Results show both RAG performance AND accuracy metrics")
console.print("4. Reports include LLM evaluation scores")
def demonstrate_evaluation_output():
"""Show what the evaluation output looks like"""
console.print("\n[bold cyan]Sample Evaluation Output[/bold cyan]")
console.print("="*80)
sample_output = """
============================================================
Running LLM Evaluation...
------------------------------------------------------------
LLM Evaluation Reward: 0.850/1.000
Passed: Yes
Reasoning: The agent correctly recalled the account number from the conversation history. The response is accurate and directly answers the user's question.
Required Information Found:
✓ account number: 123456789
============================================================
============================================================
Evaluation Complete for demo_test
LLM Evaluation Passed: ✓
LLM Reward Score: 0.850/1.000
Iterations: 2
Tool Calls: 3
Chunks: 1
Processing Time: 1.23s
Indexing Time: 0.45s
============================================================
"""
console.print(Panel(sample_output, title="Expected Console Output", border_style="green"))
console.print("\n[bold]Key Features:[/bold]")
console.print("• Automatic evaluation after each test")
console.print("• Continuous reward score (0.0 to 1.0)")
console.print("• Pass/fail determination (>= 0.6 passes)")
console.print("• Detailed reasoning for the score")
console.print("• Verification of required information")
console.print("• Integration with existing metrics")
def check_dependencies():
"""Check if all dependencies are available"""
console.print("\n[bold cyan]Checking Dependencies[/bold cyan]")
console.print("="*80)
# Check if user-memory-evaluation is accessible
eval_path = Path(__file__).parent.parent.parent / "week2" / "user-memory-evaluation"
if eval_path.exists():
console.print(f"[green]✓ user-memory-evaluation found at: {eval_path}[/green]")
# Check for required files
required_files = [
"evaluator.py",
"models.py",
"config.py"
]
for file in required_files:
if (eval_path / file).exists():
console.print(f" [green]✓ {file}[/green]")
else:
console.print(f" [red]✗ {file} missing[/red]")
else:
console.print(f"[red]✗ user-memory-evaluation not found at: {eval_path}[/red]")
console.print("[yellow] LLM evaluation will not be available[/yellow]")
# Check for API keys
console.print("\n[bold]API Keys:[/bold]")
api_keys = {
"OPENAI_API_KEY": "OpenAI (for LLM evaluation)",
"KIMI_API_KEY": "Kimi (for agent responses)"
}
for key, description in api_keys.items():
if os.getenv(key):
console.print(f" [green]✓ {key} set ({description})[/green]")
else:
console.print(f" [yellow]⚠ {key} not set ({description})[/yellow]")
def main():
"""Run all tests"""
console.print(Panel.fit(
"[bold]LLM Evaluation Integration Test[/bold]\n"
"Verifying automatic evaluation of agent responses",
border_style="cyan"
))
try:
# Check dependencies
check_dependencies()
# Test integration
test_llm_evaluation_integration()
# Show sample output
demonstrate_evaluation_output()
console.print("\n" + "="*80)
console.print("[bold green]Integration Summary:[/bold green]")
console.print("✓ LLM evaluation is integrated into the evaluation pipeline")
console.print("✓ Results are automatically evaluated after agent responses")
console.print("✓ Evaluation metrics are displayed and saved")
console.print("✓ Reports include LLM evaluation scores")
console.print("\n[yellow]Note: Actual LLM evaluation requires valid API keys[/yellow]")
console.print("="*80 + "\n")
except Exception as e:
console.print(f"\n[red]Error during testing: {e}[/red]")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()
test_logging.py¶
#!/usr/bin/env python3
"""Test script to verify tool logging and full content retrieval"""
import os
import json
from rich.console import Console
from chunker import ConversationChunker, ConversationChunk, ConversationMessage
from indexer import MemoryIndexer
from tools import MemoryTools
from config import Config
# Set dummy API key for testing
os.environ["KIMI_API_KEY"] = "test_key"
console = Console()
def create_test_chunks():
"""Create test conversation chunks with substantial content"""
chunks = []
# Create a test chunk with multiple messages
messages = []
for i in range(20):
messages.append(ConversationMessage(
role="user",
content=f"User message {i+1}: This is a detailed question about topic {i+1}. " +
f"I need information about account #{1000000 + i} and transaction date 2024-{(i%12)+1:02d}-{(i%28)+1:02d}."
))
messages.append(ConversationMessage(
role="assistant",
content=f"Assistant response {i+1}: Regarding your question about topic {i+1}, " +
f"your account #{1000000 + i} shows a transaction on 2024-{(i%12)+1:02d}-{(i%28)+1:02d} " +
f"with amount ${'1' * ((i%3)+1) + '00.00'}. Additional details include reference code REF{2000+i}."
))
chunk = ConversationChunk(
chunk_id="test_chunk_full_content",
conversation_id="conv_test_001",
test_id="test_logging",
chunk_index=0,
start_round=1,
end_round=20,
messages=messages,
metadata={
"business": "Test Bank",
"department": "Customer Service",
"agent": "Test Agent",
"duration": "45 minutes"
},
context_before="Previous conversation discussed account opening procedures",
context_after="Next conversation will cover credit card applications"
)
chunks.append(chunk)
# Create a second chunk for context testing
chunk2 = ConversationChunk(
chunk_id="test_chunk_context",
conversation_id="conv_test_001",
test_id="test_logging",
chunk_index=1,
start_round=21,
end_round=25,
messages=[
ConversationMessage(role="user", content="What about my credit card application?"),
ConversationMessage(role="assistant", content="Your credit card application #CC2024001 is approved.")
],
metadata={"business": "Test Bank"}
)
chunks.append(chunk2)
return chunks
def test_search_memory_logging():
"""Test search_memory tool with full logging"""
console.print("\n[bold cyan]Testing search_memory with full logging[/bold cyan]")
console.print("="*60)
# Initialize components
config = Config.from_env()
indexer = MemoryIndexer(config.index)
# Add test chunks
chunks = create_test_chunks()
indexer.add_chunks(chunks, rebuild=False) # Don't send to pipeline for this test
# Initialize tools
tools = MemoryTools(indexer)
# Test search with detailed query
console.print("\n[yellow]Executing search_memory...[/yellow]")
result = tools.search_memory(
query="account number transaction date reference code",
top_k=3
)
# Display results
console.print("\n[green]Search Results:[/green]")
if result.success:
data = result.data
console.print(f"Total results: {data['total_results']}")
for idx, res in enumerate(data['results'], 1):
console.print(f"\n[cyan]Result {idx}:[/cyan]")
console.print(f" Chunk ID: {res['chunk_id']}")
console.print(f" Score: {res['score']}")
console.print(f" Content length: {len(res['content'])} characters")
# Show first 500 chars to verify it's not truncated
console.print(f" Content preview (first 500 chars):")
console.print(f" {res['content'][:500]}...")
# Verify full content is present
if len(res['content']) > 1000:
console.print(f" [green]✓ Full content retrieved ({len(res['content'])} chars)[/green]")
else:
console.print(f" [yellow]⚠ Content might be truncated ({len(res['content'])} chars)[/yellow]")
def test_get_conversation_context_logging():
"""Test get_conversation_context with full logging"""
console.print("\n[bold cyan]Testing get_conversation_context with full logging[/bold cyan]")
console.print("="*60)
# Initialize components
config = Config.from_env()
indexer = MemoryIndexer(config.index)
# Add test chunks
chunks = create_test_chunks()
indexer.add_chunks(chunks, rebuild=False)
# Initialize tools
tools = MemoryTools(indexer)
# Test getting context
console.print("\n[yellow]Executing get_conversation_context...[/yellow]")
result = tools.get_conversation_context(
chunk_id="test_chunk_full_content",
context_size=2
)
# Display results
console.print("\n[green]Context Results:[/green]")
if result.success:
data = result.data
# Check target chunk
target = data['target_chunk']
console.print(f"\n[cyan]Target Chunk:[/cyan]")
console.print(f" Chunk ID: {target['chunk_id']}")
console.print(f" Rounds: {target['rounds']}")
console.print(f" Content length: {len(target['content'])} characters")
if len(target['content']) > 1000:
console.print(f" [green]✓ Full content retrieved ({len(target['content'])} chars)[/green]")
# Check context chunks
console.print(f"\n[cyan]Context Chunks: {len(data['context_chunks'])}[/cyan]")
for ctx in data['context_chunks']:
console.print(f" • {ctx['chunk_id']} ({ctx['position']}): {len(ctx['content'])} chars")
if len(ctx['content']) > 100:
console.print(f" [green]✓ Full context content[/green]")
def test_get_full_conversation_logging():
"""Test get_full_conversation with full logging"""
console.print("\n[bold cyan]Testing get_full_conversation with full logging[/bold cyan]")
console.print("="*60)
# Initialize components
config = Config.from_env()
indexer = MemoryIndexer(config.index)
# Add test chunks
chunks = create_test_chunks()
indexer.add_chunks(chunks, rebuild=False)
# Initialize tools
tools = MemoryTools(indexer)
# Test getting full conversation
console.print("\n[yellow]Executing get_full_conversation...[/yellow]")
result = tools.get_full_conversation(
conversation_id="conv_test_001",
test_id="test_logging"
)
# Display results
console.print("\n[green]Full Conversation Results:[/green]")
if result.success:
data = result.data
console.print(f"Conversation ID: {data['conversation_id']}")
console.print(f"Total chunks: {data['total_chunks']}")
console.print(f"Total rounds: {data['total_rounds']}")
total_content_length = 0
for chunk in data['chunks']:
content_length = len(chunk['content'])
total_content_length += content_length
console.print(f"\n Chunk {chunk['chunk_index']}: {content_length} characters")
if content_length > 1000:
console.print(f" [green]✓ Full chunk content[/green]")
console.print(f"\n[green]Total content retrieved: {total_content_length} characters[/green]")
def test_tool_definitions():
"""Verify extract_key_information is removed from tool definitions"""
console.print("\n[bold cyan]Verifying tool definitions[/bold cyan]")
console.print("="*60)
from tools import get_tool_definitions
tools = get_tool_definitions()
tool_names = [t['function']['name'] for t in tools]
console.print("\n[green]Available tools:[/green]")
for name in tool_names:
console.print(f" • {name}")
if "extract_key_information" in tool_names:
console.print("\n[red]✗ extract_key_information still present![/red]")
else:
console.print("\n[green]✓ extract_key_information successfully removed[/green]")
console.print(f"\n[green]Total tools available: {len(tools)}[/green]")
def main():
"""Run all logging tests"""
console.print("\n[bold]Testing Tool Logging and Full Content Retrieval[/bold]")
console.print("="*80)
try:
# Test each tool
test_search_memory_logging()
test_get_conversation_context_logging()
test_get_full_conversation_logging()
test_tool_definitions()
console.print("\n" + "="*80)
console.print("[bold green]✓ All tests completed successfully![/bold green]")
console.print("\nKey validations:")
console.print(" • Tool calls are logged with full parameters")
console.print(" • Tool results are logged completely")
console.print(" • Content is NOT truncated in results")
console.print(" • extract_key_information tool is removed")
console.print("="*80 + "\n")
except Exception as e:
console.print(f"\n[red]Error during testing: {e}[/red]")
import traceback
traceback.print_exc()
if __name__ == "__main__":
main()
test_memory_card_generation.py¶
test_pipeline.py¶
#!/usr/bin/env python3
"""Test script to verify the retrieval pipeline integration"""
import requests
import json
from rich.console import Console
console = Console()
def test_retrieval_pipeline():
"""Test if retrieval pipeline is available and working"""
console.print("[bold]Testing Retrieval Pipeline Connection[/bold]\n")
# Test health endpoint
try:
response = requests.get("http://localhost:4242/health", timeout=2)
if response.status_code == 200:
console.print("[green]✓ Retrieval pipeline is available on port 4242[/green]")
else:
console.print(f"[yellow]⚠ Pipeline returned status {response.status_code}[/yellow]")
return False
except requests.exceptions.RequestException as e:
console.print("[red]✗ Retrieval pipeline not available[/red]")
console.print(f"Error: {e}")
console.print("\nTo start the retrieval pipeline:")
console.print("[cyan]cd projects/week3/retrieval-pipeline[/cyan]")
console.print("[cyan]python api_server.py[/cyan]")
return False
# Test indexing endpoint
console.print("\n[bold]Testing Indexing Capability[/bold]")
test_doc = {
"text": "This is a test document for the retrieval pipeline.",
"metadata": {
"doc_id": "test_doc_1",
"type": "test"
}
}
try:
# Clear existing index
requests.post("http://localhost:4242/clear")
# Index test document (single document format)
response = requests.post(
"http://localhost:4242/index",
json=test_doc
)
if response.status_code == 200:
result = response.json()
console.print(f"[green]✓ Successfully indexed document with ID {result.get('doc_id', 'unknown')}[/green]")
else:
console.print(f"[yellow]⚠ Indexing returned status {response.status_code}[/yellow]")
except requests.exceptions.RequestException as e:
console.print(f"[red]✗ Error indexing documents: {e}[/red]")
return False
# Test search endpoint
console.print("\n[bold]Testing Search Capability[/bold]")
try:
response = requests.post(
"http://localhost:4242/search",
json={
"query": "test document",
"mode": "hybrid",
"top_k": 5
}
)
if response.status_code == 200:
result = response.json()
# Check for any type of results
has_results = any([
result.get("dense_results"),
result.get("sparse_results"),
result.get("reranked_results")
])
if has_results:
console.print("[green]✓ Search endpoint working correctly[/green]")
else:
console.print("[yellow]⚠ Search returned no results (index might be empty)[/yellow]")
else:
console.print(f"[yellow]⚠ Search returned status {response.status_code}[/yellow]")
except requests.exceptions.RequestException as e:
console.print(f"[red]✗ Error searching: {e}[/red]")
return False
console.print("\n[green]✓ All retrieval pipeline tests passed![/green]")
return True
def test_indexer():
"""Test the modified indexer"""
console.print("\n[bold]Testing Modified Indexer[/bold]\n")
try:
from config import Config, IndexConfig
from indexer import MemoryIndexer
from chunker import ConversationChunk, ConversationMessage
# Create test configuration
config = IndexConfig()
# Initialize indexer
indexer = MemoryIndexer(config)
console.print("[green]✓ Indexer initialized successfully[/green]")
# Create a test chunk
test_chunk = ConversationChunk(
chunk_id="test_chunk_1",
conversation_id="conv_1",
test_id="test_1",
chunk_index=0,
start_round=1,
end_round=5,
messages=[
ConversationMessage(role="user", content="Hello"),
ConversationMessage(role="assistant", content="Hi there!")
],
metadata={"test": True}
)
# Add chunk to indexer
indexer.add_chunks([test_chunk])
console.print("[green]✓ Successfully added test chunk to indexer[/green]")
# Test search
results = indexer.search("hello", top_k=5)
if results:
console.print(f"[green]✓ Search returned {len(results)} result(s)[/green]")
else:
console.print("[yellow]⚠ Search returned no results[/yellow]")
console.print("\n[green]✓ Indexer tests completed![/green]")
return True
except Exception as e:
console.print(f"[red]✗ Error testing indexer: {e}[/red]")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
console.print("[bold cyan]Testing Agentic RAG for User Memory - Pipeline Integration[/bold cyan]\n")
# Test retrieval pipeline
pipeline_ok = test_retrieval_pipeline()
# Test indexer if pipeline is available
if pipeline_ok:
indexer_ok = test_indexer()
if indexer_ok:
console.print("\n[bold green]All tests passed! The system is ready to use.[/bold green]")
else:
console.print("\n[bold yellow]Indexer needs attention, but pipeline is working.[/bold yellow]")
else:
console.print("\n[bold red]Please start the retrieval pipeline first![/bold red]")
console.print("Run the following commands in a separate terminal:")
console.print("[cyan]cd /Users/boj/ai-agent-book/projects/week3/retrieval-pipeline[/cyan]")
console.print("[cyan]python api_server.py[/cyan]")
test_proactive_service.py¶
#!/usr/bin/env python3
"""Test script to demonstrate the agent's proactive service (主动服务) capabilities"""
import logging
from datetime import datetime, timedelta
from contextual_indexer import ContextualMemoryIndexer
from contextual_agent import ContextualUserMemoryAgent
from advanced_memory_manager import AdvancedMemoryCard
from config import Config
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def test_proactive_service():
"""Test the agent's ability to provide proactive service"""
print("\n" + "="*80)
print("测试主动服务 (Testing Proactive Service)")
print("="*80)
# Initialize system
config = Config.from_env()
user_id = "proactive_test_user"
indexer = ContextualMemoryIndexer(
user_id=user_id,
index_config=config.index,
chunking_config=config.chunking,
use_contextual=False
)
# Add test memory cards with potential issues
current_date = datetime.now()
# 1. Passport expiring soon
passport_card = AdvancedMemoryCard(
category="travel",
card_key="passport_info",
backstory="User mentioned passport details when booking international travel",
date_created=current_date.strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"passport_number": "XXXXX1234",
"expiration_date": (current_date + timedelta(days=45)).strftime('%Y-%m-%d'),
"issuing_country": "USA"
}
)
indexer.memory_manager.add_card(passport_card)
# 2. Upcoming travel plan
travel_card = AdvancedMemoryCard(
category="travel",
card_key="tokyo_trip_jan_2025",
backstory="User booked a trip to Tokyo for late January",
date_created=current_date.strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"destination": "Tokyo, Japan",
"departure_date": (current_date + timedelta(days=30)).strftime('%Y-%m-%d'),
"return_date": (current_date + timedelta(days=37)).strftime('%Y-%m-%d'),
"airline": "United Airlines",
"booking_reference": "UA1234567"
}
)
indexer.memory_manager.add_card(travel_card)
# 3. Medical appointment
medical_card = AdvancedMemoryCard(
category="medical",
card_key="annual_checkup_2025",
backstory="User scheduled annual physical exam",
date_created=current_date.strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"appointment_type": "Annual Physical",
"doctor": "Dr. Sarah Chen",
"clinic": "Portland Medical Center",
"date": (current_date + timedelta(days=5)).strftime('%Y-%m-%d'),
"time": "09:00 AM",
"fasting_required": True
}
)
indexer.memory_manager.add_card(medical_card)
# 4. Insurance card
insurance_card = AdvancedMemoryCard(
category="insurance",
card_key="travel_insurance_2024",
backstory="User has annual travel insurance that needs renewal",
date_created=current_date.strftime('%Y-%m-%d %H:%M:%S'),
person="Jessica Thompson (primary)",
relationship="primary account holder",
data={
"provider": "SafeTravel Insurance",
"policy_number": "ST-2024-789456",
"expiration_date": (current_date + timedelta(days=20)).strftime('%Y-%m-%d'),
"coverage": "International travel medical and trip cancellation"
}
)
indexer.memory_manager.add_card(insurance_card)
# Initialize agent
agent = ContextualUserMemoryAgent(
indexer=indexer,
config=config
)
print("\n" + "="*80)
print("Scenario: User asks about Tokyo trip preparation")
print("Expected: Agent should proactively identify passport expiration risk")
print("="*80)
# Test questions that should trigger proactive service
test_questions = [
"我一月底的东京之行，还有什么要准备的吗？",
"What do I need for my Tokyo trip?",
"我下周有什么安排吗？",
]
for i, question in enumerate(test_questions, 1):
print(f"\n{'='*60}")
print(f"Test {i}: {question}")
print('='*60)
trajectory = agent.answer_question(
question=question,
test_id=f"proactive_test_{i}",
max_iterations=5,
stream=False
)
print("\n📝 Agent Response:")
print("-" * 40)
print(trajectory.final_answer)
print("-" * 40)
# Check if agent identified key issues
if trajectory.final_answer:
answer_lower = trajectory.final_answer.lower()
print("\n✅ Proactive Service Check:")
# Check if passport expiration was mentioned
if "passport" in answer_lower and ("expir" in answer_lower or "过期" in answer_lower):
print(" ✓ Identified passport expiration risk")
else:
print(" ✗ Missed passport expiration risk")
# Check if insurance was mentioned
if "insurance" in answer_lower or "保险" in answer_lower:
print(" ✓ Mentioned travel insurance status")
else:
print(" ✗ Missed insurance consideration")
# Check if medical appointment was mentioned (for weekly schedule question)
if i == 3 and ("appointment" in answer_lower or "physical" in answer_lower or "医生" in answer_lower):
print(" ✓ Reminded about medical appointment")
# Check for urgency markers
if any(marker in trajectory.final_answer for marker in ["⚠️", "🔴", "⏰", "需要立即", "urgent", "ASAP"]):
print(" ✓ Used urgency markers for time-sensitive items")
print(f"\nMemory Cards Used: {trajectory.memory_cards_used}")
print(f"Iterations: {len(trajectory.iterations)}")
print("\n" + "="*80)
print("主动服务测试完成 (Proactive Service Test Complete)")
print("="*80)
print("\nKey Features Demonstrated:")
print("1. Risk Detection: Identifying passport expiration before travel")
print("2. Comprehensive Assistance: Connecting travel with insurance needs")
print("3. Proactive Reminders: Highlighting upcoming appointments")
print("4. Urgency Indicators: Using markers for time-sensitive matters")
if __name__ == "__main__":
test_proactive_service()
test_quick.py¶
#!/usr/bin/env python3
"""Quick test to verify the fixes work"""
from config import Config
from contextual_evaluator import ContextualMemoryEvaluator
# Test configuration loading
config = Config.from_env()
print(f"✓ Config loaded successfully")
print(f" - EvaluationConfig.use_llm_judge: {config.evaluation.use_llm_judge}")
print(f" - AgentConfig.verbose: {config.agent.verbose}")
# Test evaluator initialization
evaluator = ContextualMemoryEvaluator(config)
print(f"✓ Evaluator initialized successfully")
# Load test cases
test_cases = evaluator.load_test_cases("layer1")
print(f"✓ Loaded {len(test_cases)} test cases")
if test_cases:
test_id = test_cases[0]
print(f"\nWill test with: {test_id}")
print("Attempting evaluation...")
try:
result = evaluator.evaluate_test_case(test_id)
print(f"✓ Evaluation completed!")
print(f" Success: {result.success}")
print(f" Iterations: {result.iterations}")
print(f" Processing time: {result.processing_time:.2f}s")
except Exception as e:
print(f"✗ Evaluation failed: {e}")
import traceback
traceback.print_exc()
else:
print("No test cases found")
test_startup.py¶
#!/usr/bin/env python3
"""Simple startup test to verify the system can initialize"""
import os
# Set a dummy API key for testing initialization
os.environ["KIMI_API_KEY"] = "test_key"
from rich.console import Console
console = Console()
def test_imports():
"""Test all imports work"""
try:
from config import Config
from chunker import ConversationChunker
from indexer import MemoryIndexer
from tools import MemoryTools
from agent import UserMemoryRAGAgent
from evaluator import UserMemoryEvaluator
console.print("[green]✓ All imports successful[/green]")
return True
except Exception as e:
console.print(f"[red]✗ Import error: {e}[/red]")
return False
def test_initialization():
"""Test component initialization"""
try:
from config import Config
from chunker import ConversationChunker
from indexer import MemoryIndexer
config = Config.from_env()
console.print("[green]✓ Config loaded[/green]")
chunker = ConversationChunker(config.chunking)
console.print("[green]✓ Chunker initialized[/green]")
indexer = MemoryIndexer(config.index)
console.print("[green]✓ Indexer initialized[/green]")
return True
except Exception as e:
console.print(f"[red]✗ Initialization error: {e}[/red]")
import traceback
traceback.print_exc()
return False
if __name__ == "__main__":
console.print("[bold]Testing Agentic RAG for User Memory - Startup[/bold]\n")
imports_ok = test_imports()
if imports_ok:
init_ok = test_initialization()
if init_ok:
console.print("\n[bold green]✓ System startup successful![/bold green]")
console.print("\nThe system is ready to use. To run the full demo:")
console.print("1. Ensure the retrieval pipeline is running on port 4242")
console.print("2. Set your API keys in .env file")
console.print("3. Run: python main.py")
else:
console.print("\n[bold yellow]⚠ Initialization needs attention[/bold yellow]")
else:
console.print("\n[bold red]✗ Import errors need to be fixed[/bold red]")
test_top_k.py¶
#!/usr/bin/env python3
"""Test that top_k parameter works correctly with the retrieval pipeline"""
import os
import sys
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
# Set dummy API key
os.environ["KIMI_API_KEY"] = "sk-test"
from config import Config
from indexer import MemoryIndexer
from chunker import ConversationChunker, ConversationChunk
def test_top_k():
"""Test that different top_k values return the correct number of results"""
config = Config.from_env()
indexer = MemoryIndexer(config)
# Create some test chunks
test_chunks = []
for i in range(10):
chunk = ConversationChunk(
chunk_id=f"test_chunk_{i}",
test_id="test_id",
conversation_id=f"conv_{i}",
messages=[
{"role": "user", "content": f"Test message {i} about banking"},
{"role": "assistant", "content": f"Response {i} about account"}
],
start_round=i*2,
end_round=(i+1)*2,
metadata={"test": f"chunk_{i}"}
)
test_chunks.append(chunk)
# Build indexes
print("Building indexes with 10 test chunks...")
indexer.build_indexes(test_chunks)
# Test different top_k values
test_values = [1, 3, 5, 10, 15]
for top_k in test_values:
print(f"\nTesting top_k={top_k}...")
results = indexer.search("banking account", top_k=top_k)
actual_count = len(results)
# The actual count should match requested top_k (up to available documents)
expected_count = min(top_k, 10) # We only have 10 chunks
if actual_count == expected_count:
print(f"✓ Correct: Requested {top_k}, got {actual_count} results")
else:
print(f"✗ Error: Requested {top_k}, got {actual_count} results (expected {expected_count})")
# Show the result IDs
if results:
result_ids = [r.chunk.chunk_id for r in results[:3]] # Show first 3
print(f" First results: {result_ids}")
print("\n" + "="*60)
print("✓ top_k parameter is now working correctly!")
print(" - The pipeline respects the requested number of results")
print(" - It retrieves more candidates initially for better reranking")
print("="*60)
if __name__ == "__main__":
# Check if retrieval pipeline is running
import requests
try:
response = requests.get("http://localhost:4242/")
print("✓ Retrieval pipeline is running")
except:
print("✗ Retrieval pipeline is not running on port 4242")
print(" Please start it with: cd projects/week3/retrieval-pipeline && python main.py")
sys.exit(1)
test_top_k()
tools.py¶
"""Tool definitions for the User Memory RAG Agent
This module provides tool definitions and implementations for searching
and retrieving information from indexed conversation memories.
"""
import json
import logging
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from indexer import MemoryIndexer, SearchResult
from config import IndexConfig
logger = logging.getLogger(__name__)
@dataclass
class ToolResult:
"""Result from a tool execution"""
success: bool
data: Any
error: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
if self.success:
return {"status": "success", "data": self.data}
else:
return {"status": "error", "error": self.error}
class MemoryTools:
"""Tools for searching and retrieving user memory information"""
def __init__(self, indexer: MemoryIndexer):
"""
Initialize memory tools
Args:
indexer: The memory indexer instance
"""
self.indexer = indexer
logger.info("Initialized memory tools")
def search_memory(self,
query: str,
top_k: int = 3,
filter_test_id: Optional[str] = None) -> ToolResult:
"""
Search user memory for relevant information
Args:
query: Natural language search query
top_k: Number of results to return
filter_test_id: Optional test ID to filter results
Returns:
ToolResult with search results
"""
try:
# Perform search
results = self.indexer.search(query, top_k=top_k)
# Filter by test ID if specified
if filter_test_id:
results = [r for r in results if r.chunk.test_id == filter_test_id]
# Format results
formatted_results = []
for result in results:
# Extract key information from the chunk
chunk_info = {
"chunk_id": result.chunk_id,
"score": round(result.score, 4),
"test_id": result.chunk.test_id,
"conversation_id": result.chunk.conversation_id,
"rounds": f"{result.chunk.start_round}-{result.chunk.end_round}",
"metadata": result.chunk.metadata,
"content": result.chunk.to_text(), # FULL content, not truncated
"match_type": result.match_type
}
formatted_results.append(chunk_info)
logger.info(f"Search query: '{query}' returned {len(formatted_results)} results")
return ToolResult(
success=True,
data={
"query": query,
"total_results": len(formatted_results),
"results": formatted_results
}
)
except Exception as e:
logger.error(f"Error in search_memory: {e}")
return ToolResult(
success=False,
data=None,
error=str(e)
)
def get_conversation_context(self,
chunk_id: str,
context_size: int = 2) -> ToolResult:
"""
Get surrounding context for a specific chunk
Args:
chunk_id: The chunk ID to get context for
context_size: Number of chunks before/after to include
Returns:
ToolResult with conversation context
"""
try:
# Get the target chunk
if chunk_id not in self.indexer.chunks:
return ToolResult(
success=False,
data=None,
error=f"Chunk {chunk_id} not found"
)
target_chunk = self.indexer.chunks[chunk_id]
# Find related chunks from same conversation
related_chunks = []
for cid, chunk in self.indexer.chunks.items():
if (chunk.conversation_id == target_chunk.conversation_id and
chunk.test_id == target_chunk.test_id):
related_chunks.append(chunk)
# Sort by chunk index
related_chunks.sort(key=lambda x: x.chunk_index)
# Find target index
target_idx = next(
(i for i, c in enumerate(related_chunks) if c.chunk_id == chunk_id),
None
)
if target_idx is None:
return ToolResult(
success=False,
data=None,
error="Could not locate chunk in conversation"
)
# Get context chunks
start_idx = max(0, target_idx - context_size)
end_idx = min(len(related_chunks), target_idx + context_size + 1)
context_chunks = related_chunks[start_idx:end_idx]
# Format result
context_data = {
"target_chunk": {
"chunk_id": target_chunk.chunk_id,
"rounds": f"{target_chunk.start_round}-{target_chunk.end_round}",
"content": target_chunk.to_text()
},
"context_chunks": []
}
for chunk in context_chunks:
if chunk.chunk_id != chunk_id:
context_data["context_chunks"].append({
"chunk_id": chunk.chunk_id,
"rounds": f"{chunk.start_round}-{chunk.end_round}",
"position": "before" if chunk.chunk_index < target_chunk.chunk_index else "after",
"content": chunk.to_text()
})
return ToolResult(
success=True,
data=context_data
)
except Exception as e:
logger.error(f"Error in get_conversation_context: {e}")
return ToolResult(
success=False,
data=None,
error=str(e)
)
def get_full_conversation(self,
conversation_id: str,
test_id: str) -> ToolResult:
"""
Retrieve all chunks from a specific conversation
Args:
conversation_id: Conversation identifier
test_id: Test case identifier
Returns:
ToolResult with full conversation
"""
try:
# Find all chunks for this conversation
conversation_chunks = []
for chunk_id, chunk in self.indexer.chunks.items():
if (chunk.conversation_id == conversation_id and
chunk.test_id == test_id):
conversation_chunks.append(chunk)
if not conversation_chunks:
return ToolResult(
success=False,
data=None,
error=f"No chunks found for conversation {conversation_id}"
)
# Sort by chunk index
conversation_chunks.sort(key=lambda x: x.chunk_index)
# Format result
conversation_data = {
"conversation_id": conversation_id,
"test_id": test_id,
"total_chunks": len(conversation_chunks),
"total_rounds": max(c.end_round for c in conversation_chunks),
"metadata": conversation_chunks[0].metadata if conversation_chunks else {},
"chunks": []
}
for chunk in conversation_chunks:
conversation_data["chunks"].append({
"chunk_id": chunk.chunk_id,
"chunk_index": chunk.chunk_index,
"rounds": f"{chunk.start_round}-{chunk.end_round}",
"content": chunk.to_text()
})
return ToolResult(
success=True,
data=conversation_data
)
except Exception as e:
logger.error(f"Error in get_full_conversation: {e}")
return ToolResult(
success=False,
data=None,
error=str(e)
)
def get_tool_definitions() -> List[Dict[str, Any]]:
"""
Get OpenAI function calling tool definitions
Returns:
List of tool definitions for OpenAI API
"""
return [
{
"type": "function",
"function": {
"name": "search_memory",
"description": "Search user conversation memory for relevant information. Use this to find specific details from past conversations.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Natural language search query describing what information to find"
},
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "get_conversation_context",
"description": "Get surrounding context for a specific conversation chunk. Use this when you need more context around a search result.",
"parameters": {
"type": "object",
"properties": {
"chunk_id": {
"type": "string",
"description": "The chunk ID to get context for"
},
"context_size": {
"type": "integer",
"description": "Number of chunks before/after to include (default: 2)",
"default": 2
}
},
"required": ["chunk_id"]
}
}
},
{
"type": "function",
"function": {
"name": "get_full_conversation",
"description": "Retrieve all chunks from a specific conversation. Use this when you need to review an entire conversation history.",
"parameters": {
"type": "object",
"properties": {
"conversation_id": {
"type": "string",
"description": "The conversation identifier"
},
"test_id": {
"type": "string",
"description": "The test case identifier"
}
},
"required": ["conversation_id", "test_id"]
}
}
}
]
verify_test_loading.py¶
#!/usr/bin/env python3
"""Verify test case loading with correct field mapping"""
import os
import sys
from pathlib import Path
# Set dummy API key
os.environ["KIMI_API_KEY"] = os.getenv("KIMI_API_KEY", "sk-test")
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-test")
from config import Config
from evaluator import UserMemoryEvaluator
# Initialize evaluator
config = Config.from_env()
evaluator = UserMemoryEvaluator(config)
# Load test cases
test_cases = evaluator.load_test_cases(category="layer1")
if test_cases:
print(f"\nLoaded {len(test_cases)} test cases")
# Check first test case
test_id = test_cases[0]
test_case = evaluator.test_cases[test_id]
print(f"\nTest Case: {test_id}")
print(f" Title: {test_case.title}")
print(f" Category: {test_case.category}")
print(f" User Question: {test_case.user_question}")
print(f" Evaluation Criteria: {test_case.evaluation_criteria[:200]}..." if test_case.evaluation_criteria else " Evaluation Criteria: None")
print(f" Expected Behavior: {test_case.expected_behavior[:100]}..." if test_case.expected_behavior else " Expected Behavior: None")
print(f" Conversations: {len(test_case.conversation_histories)}")
# Check conversation history structure
if test_case.conversation_histories:
conv = test_case.conversation_histories[0]
print(f"\nFirst Conversation:")
print(f" ID: {conv.get('conversation_id', 'N/A')}")
print(f" Timestamp: {conv.get('timestamp', 'N/A')}")
print(f" Messages: {len(conv.get('messages', []))}")
print(f" Has metadata: {bool(conv.get('metadata', {}))}")
print("\n✓ Test case structure is correct")
else:
print("✗ No test cases loaded")
memory_qa_eval.json¶
{
"description": "受控的用户记忆问答对比集（实验 3-12）。用于离线量化『在嵌入/索引前对记忆块做上下文化』相较『直接索引原始对话块』在召回用户事实上的提升。每个 chunk 同时给出 text（原始对话块）与 context（上下文前缀，代表由 LLM 生成的『情境锚定』摘要）。gold 为该 query 对应的正确记忆块 id。数据集刻意包含若干『孤立片段』（如『好的，就订这个吧』），它们脱离上下文后几乎不含可检索信号——这正是本章要说明的痛点。",
"note": "本文件为教学用受控数据集；生产环境的 context 由 LLM 逐块生成、并用神经嵌入检索（需 API）。此处以确定性的 BM25 词法检索作为无需 API 的离线代理，度量的是同一份 context 分别拼接前后对召回的影响。",
"chunks": [
{
"id": "flight_detail",
"scenario": "flight",
"text": "用户：帮我看看从上海到西雅图的单程航班。\n助手：为您找到达美航空 DL588，10 月 12 日出发，单程含税 500 美元。",
"context": "本段来自用户与旅行助手 2024 年 10 月的机票预订对话，讨论上海飞西雅图的单程航班：达美 DL588，含税 500 美元。"
},
{
"id": "flight_confirm",
"scenario": "flight",
"text": "用户：好的，就订这个吧。\n助手：已为您预订，稍后发送电子行程单。",
"context": "本段承接上文，用户确认预订那张 500 美元、上海飞西雅图的达美 DL588 单程机票。"
},
{
"id": "hotel_detail",
"scenario": "hotel",
"text": "用户：西雅图市中心有什么酒店推荐？\n助手：推荐凯悦酒店，10 月 12 日至 15 日三晚共 840 美元。",
"context": "本段来自同一次西雅图行程的酒店预订讨论：凯悦酒店，10 月 12-15 日三晚 840 美元。"
},
{
"id": "hotel_confirm",
"scenario": "hotel",
"text": "用户：可以，帮我订下来。\n助手：好的，已锁定房间。",
"context": "本段承接上文，用户确认预订西雅图凯悦酒店三晚（840 美元）的房间。"
},
{
"id": "bank_account",
"scenario": "bank",
"text": "助手：您的新支票账户已开通。\n用户：号码是多少？\n助手：账号是 4429853327，路由号 123006800。",
"context": "本段来自用户在第一国民银行开设支票账户的通话，客服告知支票账号 4429853327 与路由号 123006800。"
},
{
"id": "bank_confirm",
"scenario": "bank",
"text": "用户：记下了，就用这个。\n助手：好的，已为您设置为默认收款账户。",
"context": "本段承接上文，用户确认使用第一国民银行支票账户 4429853327 作为默认收款账户。"
},
{
"id": "passport",
"scenario": "travel_doc",
"text": "用户：顺便问下，我护照还有效吧？\n助手：您的护照 2025 年 2 月 18 日到期，出行前请留意。",
"context": "本段提到用户 Jessica 的护照将于 2025 年 2 月 18 日过期，与其一月的东京出行相关。"
},
{
"id": "medication",
"scenario": "medical",
"text": "医生：这次给你把剂量调整一下。\n用户：好的，听你的。",
"context": "本段来自复诊对话，医生将用户的二甲双胍剂量由每日 500 毫克上调至每日 1000 毫克。"
},
{
"id": "restaurant_pref",
"scenario": "preference",
"text": "用户：我不吃香菜，也对花生过敏。\n助手：已记录您的饮食禁忌。",
"context": "本段记录用户的饮食偏好与过敏信息：不吃香菜，对花生过敏。"
},
{
"id": "gym_distractor",
"scenario": "gym",
"text": "用户：我想办张健身卡。\n助手：我们有月卡和年卡两种，年卡 1200 元。",
"context": "本段来自健身房会籍咨询，年卡 1200 元。"
},
{
"id": "car_distractor",
"scenario": "car",
"text": "用户：租车大概多少钱一天？\n助手：经济型每天 45 美元，含基础保险。",
"context": "本段来自租车询价，经济型每天 45 美元。"
},
{
"id": "insurance_distractor",
"scenario": "insurance",
"text": "用户：我的车险什么时候续？\n助手：您的保单 12 月 1 日到期，可提前一个月续保。",
"context": "本段来自车险咨询，保单 12 月 1 日到期。"
}
],
"queries": [
{"q": "我最后确认预订的那张上海到西雅图的机票是哪个航班？", "gold": "flight_confirm"},
{"q": "我订的西雅图酒店确认了吗，是哪家？", "gold": "hotel_confirm"},
{"q": "我设为默认收款的银行支票账户是哪个？", "gold": "bank_confirm"},
{"q": "从上海到西雅图那趟航班的票价是多少？", "gold": "flight_detail"},
{"q": "我的护照什么时候过期？", "gold": "passport"},
{"q": "医生把我的二甲双胍剂量调成多少了？", "gold": "medication"},
{"q": "我有哪些饮食过敏和忌口？", "gold": "restaurant_pref"},
{"q": "西雅图凯悦酒店三晚一共多少钱？", "gold": "hotel_detail"}
]
}