agentic-rag-for-user-memory¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/agentic-rag-for-user-memory
项目说明¶
Agentic RAG for User Memory Evaluation¶
An educational project that combines Retrieval-Augmented Generation (RAG) with User Memory Evaluation to demonstrate how AI agents can effectively manage and query long-term conversation histories.
🎯 Learning Objectives¶
This project teaches you: 1. How to chunk long conversations into manageable segments for indexing 2. How to integrate with external retrieval pipelines for hybrid search 3. How to implement agentic RAG with tool-calling and the ReAct pattern 4. How to evaluate memory systems with automatic LLM-based scoring 5. How to optimize retrieval for conversation-based queries 6. How to integrate evaluation frameworks from different projects
🏗️ Architecture Overview¶
┌─────────────────────────────────────────┐
│ User Memory Test Cases │
│ (60 test cases, 3 difficulty layers) │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Conversation Chunker │
│ (Splits into 20-round segments with │
│ overlap and contextual enrichment) │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ External Retrieval Pipeline │
│ (Port 4242 - Hybrid Search) │
│ ┌─────────────┐ ┌──────────────────┐ │
│ │Dense Search │ │ Sparse Search │ │
│ │ (Embeddings)│ │ (BM25) │ │
│ └─────────────┘ └──────────────────┘ │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Agentic RAG Agent │
│ (ReAct pattern with memory tools) │
│ │
│ Tools: │
│ • search_memory │
│ • get_conversation_context │
│ • get_full_conversation │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ LLM Evaluation System │
│ (Automatic scoring and reasoning) │
│ • Reward score (0.0-1.0) │
│ • Pass/Fail determination │
│ • Detailed reasoning │
└─────────────────────────────────────────┘
📚 Key Concepts¶
1. Conversation Chunking¶
Long conversation histories are divided into chunks of approximately 20 rounds (user-assistant exchanges). This makes them: - Searchable: Smaller units are easier to index and retrieve - Contextual: Each chunk maintains context from surrounding conversations - Efficient: Reduces the amount of text the LLM needs to process
2. Hybrid Retrieval (via External Pipeline)¶
The system integrates with an external retrieval pipeline service that provides: - Dense Retrieval: Uses embeddings for semantic similarity search - Sparse Retrieval: Uses BM25 for keyword matching and exact phrase search - Hybrid Fusion: Combines scores from both methods for optimal results - Scalable Architecture: Offloads indexing and search to dedicated service
3. Agentic RAG Pattern¶
The agent follows the ReAct (Reasoning + Acting) pattern: 1. Reason about what information is needed 2. Act by calling search tools 3. Observe the results 4. Iterate until sufficient information is found
4. Automatic LLM Evaluation¶
The system integrates with week2/user-memory-evaluation to provide: - Reward Scoring: Continuous score from 0.0 to 1.0 - Pass/Fail Assessment: Automatic determination (≥0.6 passes) - Detailed Reasoning: Explanation of evaluation decisions - Full Visibility: Enhanced logging of LLM responses and tool calls
5. Contextual Enrichment¶
Chunks are enhanced with: - Metadata about the conversation (business, department, timestamps) - Context from previous and next chunks - Semantic tags for better retrieval
🚀 Quick Start¶
Prerequisites¶
- Python 3.8+
- Retrieval Pipeline Service on port 4242 is now OPTIONAL. By default the indexer
uses
retrieval_backend="auto": it uses the external pipeline if it is reachable, otherwise it transparently falls back to a built-in, dependency-free local BM25 index so the whole chunk → index → retrieve path runs fully offline. - API keys are only needed for the LLM-driven parts:
- A supported LLM provider (Kimi/Moonshot, OpenAI, SiliconFlow, DeepSeek, ...) for
agent answer generation (
--mode batch/interactive/demo). - The offline comparison demo (
--mode offline-demo) needs NO API key and NO port 4242 service.
Installation¶
# Enter the project
cd chapter3/agentic-rag-for-user-memory
# Install dependencies
pip install -r requirements.txt
# Setup environment variables
cp env.example .env
# Edit .env with your API keys
Retrieval Backend (offline by default, pipeline optional)¶
The retrieval backend is selectable via --backend (or IndexConfig.retrieval_backend):
| value | behavior |
|---|---|
auto |
default — use the port-4242 pipeline if reachable, else local BM25 |
local |
always use the built-in offline BM25 index (no external service) |
pipeline |
always use the external retrieval pipeline on port 4242 |
To use the external pipeline (for dense/hybrid embeddings + reranking), start it first:
# In a separate terminal (OPTIONAL)
cd ../retrieval-pipeline
python api_server.py # serves http://localhost:4242
Running the Demo¶
# Offline comparison demo — NO API key, NO port 4242 needed.
# Shows agentic multi-hop memory retrieval beating naive single-query recall,
# on the multi-session layer2_01_multiple_vehicles case, with a metric table.
python main.py --mode offline-demo
python offline_demo.py # equivalent, standalone entry point
python offline_demo.py --output results/offline_demo.json # also dump JSON
# Test the system setup
python test_pipeline.py
# Run interactive mode
python main.py
# Quick demo with a simple test case (needs an LLM API key)
python main.py --mode demo
# Batch evaluation of a category (needs an LLM API key; --backend local stays offline)
python main.py --mode batch --category layer1 --backend local
Offline demo result (reproducible)¶
Running python offline_demo.py on layer2_01_multiple_vehicles (user owns a Honda
Accord with a scheduled Firestone service and a Tesla Model 3 without one, discussed
across two separate calls) yields, from real BM25 retrieval:
| metric | naive single-query | agentic multi-hop |
|---|---|---|
| retrieval queries issued | 1 | 5 |
| memory chunks retrieved | 3 | 5 |
| decisive-evidence recall | 50% | 100% |
| can fully disambiguate & answer | no | yes |
The naive query is dominated by "schedule service" keywords and misses the Honda
appointment-confirmation chunk (FS-447291). The agentic strategy discovers the
second vehicle from the first-round results, issues focused follow-up queries per
vehicle, and recovers the missing evidence. The recall numbers are computed from
actual retrieval, not hard-coded.
CLI flags (main.py)¶
--mode {interactive,batch,demo,offline-demo}, --category, --test-id, --query,
--provider, --model, --index-mode {dense,sparse,hybrid},
--backend {auto,local,pipeline}, --top-k, --rounds-per-chunk, --store-path,
--test-cases-dir, --output, --config. Run python main.py --help for the
(Chinese) descriptions.
📖 Usage Guide¶
Interactive Mode¶
The interactive interface provides these options:
- Load Test Cases: Load test cases from the evaluation framework
- View Test Cases: Browse loaded test cases and their details
- Configure Settings: Adjust chunking, indexing, and agent parameters
- Evaluate Single Test: Run evaluation on a specific test case
- Evaluate by Category: Test all cases in a difficulty layer
- Generate Report: Create detailed evaluation reports
Example Workflow¶
# 1. Initialize the evaluator
from config import Config
from evaluator import UserMemoryEvaluator
config = Config.from_env()
evaluator = UserMemoryEvaluator(config)
# 2. Load test cases
test_cases = evaluator.load_test_cases(category="layer1")
# 3. Evaluate a test case
result = evaluator.evaluate_test_case("layer1_01_bank_account")
# 4. Generate report
report = evaluator.generate_report("results/evaluation_report.txt")
Configuring the System¶
Key configuration options in config.py:
# Chunking settings
config.chunking.rounds_per_chunk = 20 # Rounds per chunk
config.chunking.overlap_rounds = 2 # Overlapping rounds
# Index settings
config.index.mode = "hybrid" # dense, sparse, or hybrid
config.index.enable_contextual = True # Add contextual enrichment
# Agent settings
config.agent.max_search_results = 5 # Results per search
config.evaluation.max_iterations = 10 # Max ReAct iterations
🧪 Test Case Structure¶
Test cases follow the user-memory-evaluation framework format:
Test Case Fields¶
test_id: Unique identifiercategory: Difficulty layer (layer1, layer2, layer3)title: Descriptive titleconversation_histories: Past conversations to indexuser_question: The question to answerevaluation_criteria: Criteria for evaluating responsesexpected_behavior: Optional expected agent behavior
Layer 1: Simple Information Retrieval¶
- Single conversation with clear information
- Direct questions about specific details
- Example: "What is my checking account number?"
Layer 2: Multi-Conversation Correlation¶
- Multiple related conversations
- Questions requiring information synthesis
- Example: "Which of my vehicles needs service first?"
Layer 3: Complex Reasoning¶
- Hidden patterns and implicit connections
- Questions requiring deep analysis
- Example: "What urgent issues should I address before my trip?"
🔧 Component Details¶
Chunker (chunker.py)¶
- Splits conversations into fixed-size chunks
- Maintains conversation flow with overlapping rounds
- Adds contextual information to each chunk
Indexer (indexer.py)¶
- Integrates with external retrieval pipeline service
- Sends documents for indexing via HTTP API
- Manages document ID mapping
- Performs hybrid searches through the pipeline
Tools (tools.py)¶
search_memory: Main search interface with full content retrievalget_conversation_context: Retrieves surrounding chunksget_full_conversation: Gets entire conversation history Note: All tools return complete content (not truncated)
Agent (agent.py)¶
- Implements ReAct pattern with tool calling
- Manages conversation state
- Generates responses based on retrieved information
Evaluator (evaluator.py)¶
- Loads test cases from YAML files
- Manages the indexing pipeline
- Tracks evaluation metrics and results
- Integrates automatic LLM evaluation
📊 Evaluation Metrics¶
The system tracks comprehensive metrics: - Success Rate: Percentage of correctly answered questions - LLM Evaluation Score: Automatic reward score (0.0-1.0) with detailed reasoning - Iterations: Number of ReAct reasoning steps - Tool Calls: Number and types of tools used - Processing Time: Response generation time - Indexing Time: Time to build search indexes - Result Quality: Controlled by reranking with configurable top_k
🔍 Troubleshooting¶
Top-K Results Issue¶
Problem: Getting 10 results regardless of top_k setting
Solution: The retrieval pipeline uses two parameters:
- top_k: Initial retrieval count (for candidates)
- rerank_top_k: Final result count (what you actually get)
The system now correctly sets both parameters to respect your requested result count.
LLM Evaluation Not Running¶
Problem: No automatic evaluation after agent response
Solution: Ensure you have:
- Valid OpenAI API key for evaluation
- Access to week2/user-memory-evaluation module
- Proper test case format with evaluation_criteria
Retrieval Pipeline Connection¶
Problem: Cannot connect to retrieval pipeline
Solution: This is no longer fatal — with --backend auto (default) the system logs a
warning and falls back to the built-in local BM25 backend. If you specifically want the
external pipeline:
- Start the service: cd ../retrieval-pipeline && python api_server.py
- Verify it's running on http://localhost:4242
- Or force offline mode explicitly with --backend local
📄 License¶
This project is part of the AI Agent training curriculum and is intended for educational purposes.
🔗 Related Projects¶
week2/user-memory: Basic user memory systemweek2/user-memory-evaluation: Evaluation frameworkweek3/agentic-rag: Original agentic RAG implementationweek3/contextual-retrieval: Advanced retrieval techniques
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.
源代码¶
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,
top_k=self.config.agent.max_search_results,
filter_test_id=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
# Retrieval backend selection:
# "auto" -> use the port-4242 retrieval pipeline if reachable, otherwise fall back
# to a built-in, dependency-free local BM25 index (works fully offline)
# "local" -> always use the built-in local BM25 index (no external service needed)
# "pipeline" -> always use the external retrieval pipeline on port 4242
retrieval_backend: str = "auto"
retrieval_url: str = "http://localhost:4242" # External retrieval pipeline endpoint
@dataclass
class EvaluationConfig:
"""Configuration for evaluation framework"""
test_cases_dir: str = "../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
@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
@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)
if backend := os.getenv("RETRIEVAL_BACKEND"):
config.index.retrieval_backend = backend
if test_cases_dir := os.getenv("TEST_CASES_DIR"):
config.evaluation.test_cases_dir = test_cases_dir
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
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 re
import math
import json
import logging
import requests
from collections import Counter
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__)
def _tokenize(text: str) -> List[str]:
"""Lightweight tokenizer shared by the local backend (alphanumeric + CJK)."""
return re.findall(r"[a-zA-Z0-9]+|[一-鿿]", text.lower())
class LocalBM25Backend:
"""A dependency-free, in-process BM25 index.
This is the offline fallback for the external retrieval pipeline: it lets the
whole store/retrieve path run without any network service or API key, which is
what makes the experiment reproducible on a laptop. Sparse (BM25) retrieval is
the same lexical scoring the pipeline exposes as its "sparse"/"hybrid" modes.
"""
def __init__(self, k1: float = 1.5, b: float = 0.75):
self.k1 = k1
self.b = b
self.doc_ids: List[str] = []
self.doc_tokens: List[List[str]] = []
self.idf: Dict[str, float] = {}
self.avgdl: float = 0.0
self._built: bool = False
def clear(self):
self.doc_ids = []
self.doc_tokens = []
self.idf = {}
self.avgdl = 0.0
self._built = False
def add(self, doc_id: str, text: str):
self.doc_ids.append(doc_id)
self.doc_tokens.append(_tokenize(text))
self._built = False
def _build(self):
n_docs = len(self.doc_tokens)
df: Counter = Counter()
for tokens in self.doc_tokens:
for term in set(tokens):
df[term] += 1
# Standard BM25 idf with +1 smoothing so it stays non-negative.
self.idf = {
term: math.log(1 + (n_docs - freq + 0.5) / (freq + 0.5))
for term, freq in df.items()
}
self.avgdl = (sum(len(t) for t in self.doc_tokens) / n_docs) if n_docs else 0.0
self._built = True
def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
if not self._built:
self._build()
query_terms = _tokenize(query)
scored: List[Tuple[str, float]] = []
for idx, tokens in enumerate(self.doc_tokens):
tf = Counter(tokens)
dl = len(tokens)
score = 0.0
for term in query_terms:
freq = tf.get(term, 0)
if not freq:
continue
idf = self.idf.get(term, 0.0)
denom = freq + self.k1 * (1 - self.b + self.b * dl / (self.avgdl or 1))
score += idf * (freq * (self.k1 + 1)) / denom
if score > 0:
scored.append((self.doc_ids[idx], score))
scored.sort(key=lambda item: item[1], reverse=True)
return scored[:top_k]
@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 = getattr(self.config, "retrieval_url", "http://localhost:4242")
# Local, in-process fallback backend (no external service required)
self.local_backend = LocalBM25Backend()
# 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)
# Decide which backend to use: "local", "pipeline", or "auto"
backend = getattr(self.config, "retrieval_backend", "auto")
if backend == "local":
self.use_local = True
elif backend == "pipeline":
self.use_local = False
self._check_retrieval_pipeline()
else: # auto
self.use_local = not self._check_retrieval_pipeline()
if self.use_local:
logger.info("Using built-in local BM25 backend (offline mode, no port 4242 needed)")
else:
logger.info("Using external retrieval pipeline backend")
logger.info(f"Initialized indexer with mode: {self.config.mode}")
def _check_retrieval_pipeline(self) -> bool:
"""Check if the retrieval pipeline service is available. Returns True if reachable."""
try:
response = requests.get(f"{self.retrieval_url}/health", timeout=2)
if response.status_code == 200:
logger.info("✓ Retrieval pipeline service is available")
return True
logger.warning(f"Retrieval pipeline returned status {response.status_code}")
return False
except requests.exceptions.RequestException as e:
logger.warning(f"Retrieval pipeline service not available at {self.retrieval_url}: {e}")
logger.info("Note: falling back to the built-in local BM25 backend (offline).")
logger.info("To use the external pipeline instead, start it with:")
logger.info(" cd ../retrieval-pipeline && python api_server.py")
return False
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]]):
"""Index documents into the active backend (local BM25 or the external pipeline)."""
if self.use_local:
self.local_backend.clear()
for doc in documents:
chunk_id = doc.get("metadata", {}).get("doc_id")
if chunk_id:
self.local_backend.add(chunk_id, doc["text"])
self.doc_id_mapping[chunk_id] = chunk_id
logger.info(f"Indexed {len(documents)} documents into local BM25 backend")
return
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
# Offline path: score against the in-process BM25 index.
if self.use_local:
results = []
for chunk_id, score in self.local_backend.search(query, top_k=top_k):
chunk = self.chunks.get(chunk_id)
if chunk:
results.append(SearchResult(
chunk_id=chunk_id,
score=float(score),
chunk=chunk,
match_type="local_bm25"
))
logger.info(f"Search returned {len(results)} results from local BM25 backend")
return results
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 Agentic RAG User Memory Evaluation System
This script provides an interactive interface for:
1. Loading test cases from the user-memory-evaluation framework
2. Chunking and indexing conversation histories
3. Evaluating the RAG agent on selected test cases
"""
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, ChunkingStrategy, IndexMode
from evaluator import UserMemoryEvaluator
# 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 InteractiveRAGEvaluator:
"""Interactive interface for the RAG evaluation system"""
def __init__(self, config: Optional[Config] = None):
"""Initialize the interactive evaluator"""
self.config = config or Config.from_env()
self.evaluator = UserMemoryEvaluator(self.config)
self.test_cases_loaded = False
def run(self):
"""Run the interactive evaluation session"""
console.print(Panel.fit(
"[bold cyan]Agentic RAG for User Memory Evaluation[/bold cyan]\n"
"Educational Project for Learning RAG + User Memory Systems",
border_style="cyan"
))
while True:
self.show_menu()
choice = Prompt.ask(
"Select an option",
choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "0"],
default="1"
)
if choice == "1":
self.load_test_cases()
elif choice == "2":
self.view_test_cases()
elif choice == "3":
self.configure_settings()
elif choice == "4":
self.evaluate_single_test()
elif choice == "5":
self.evaluate_category()
elif choice == "6":
self.evaluate_all()
elif choice == "7":
self.view_results()
elif choice == "8":
self.generate_report()
elif choice == "9":
self.demo_mode()
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. Load Test Cases")
console.print("2. View Loaded Test Cases")
console.print("3. Configure Settings")
console.print("4. Evaluate Single Test Case")
console.print("5. Evaluate by Category")
console.print("6. Evaluate All Test Cases")
console.print("7. View Results")
console.print("8. Generate Report")
console.print("9. Demo Mode (Quick Test)")
console.print("0. Exit")
def load_test_cases(self):
"""Load test cases from the evaluation framework"""
console.print("\n[cyan]Loading test cases...[/cyan]")
category = Prompt.ask(
"Select category to load",
choices=["all", "layer1", "layer2", "layer3"],
default="all"
)
category_filter = None if category == "all" else category
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
console=console
) as progress:
task = progress.add_task("Loading test cases...", total=None)
test_cases = self.evaluator.load_test_cases(category_filter)
progress.update(task, completed=True)
self.test_cases_loaded = True
console.print(f"[green]✓ Loaded {len(test_cases)} test cases[/green]")
# Show summary
categories = {}
for test_id in test_cases:
tc = self.evaluator.test_cases[test_id]
cat = tc.category
categories[cat] = categories.get(cat, 0) + 1
for cat, count in sorted(categories.items()):
console.print(f" {cat}: {count} test cases")
def view_test_cases(self):
"""View loaded test cases"""
if not self.test_cases_loaded:
console.print("[yellow]No test cases loaded. Please load test cases first.[/yellow]")
return
# Create table
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")
# Sort test cases for consistent numbering
sorted_test_cases = sorted(self.evaluator.test_cases.items())
for idx, (test_id, test_case) in enumerate(sorted_test_cases, 1):
table.add_row(
str(idx),
test_id,
test_case.category,
test_case.title[:50] + "..." if len(test_case.title) > 50 else test_case.title,
str(len(test_case.conversation_histories))
)
console.print(table)
# Option to view details
if Confirm.ask("\nView test case details?"):
# Build a list of test IDs for selection
test_ids_list = list(self.evaluator.test_cases.keys())
test_ids_list.sort()
console.print("\n[dim]Enter test ID or its index number from the table above[/dim]")
user_input = Prompt.ask("Select test case")
# Check if user entered a number (1-based index from table display)
test_id = None
if user_input.isdigit():
idx = int(user_input) - 1
if 0 <= idx < len(test_ids_list):
test_id = test_ids_list[idx]
else:
console.print(f"[red]Invalid index number: {user_input}[/red]")
return
else:
test_id = user_input
if test_id in self.evaluator.test_cases:
test_case = self.evaluator.test_cases[test_id]
console.print(Panel(
f"[bold]Test Case: {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)}"
+ (f"\nExpected Behavior: {test_case.expected_behavior[:100]}..." if test_case.expected_behavior else ""),
title=test_id,
border_style="cyan"
))
else:
console.print(f"[red]Test case {test_id} not found[/red]")
def configure_settings(self):
"""Configure RAG and evaluation settings"""
console.print("\n[bold]Configuration Settings[/bold]")
# 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 Strategy: {self.config.chunking.strategy}")
console.print(f" Rounds per Chunk: {self.config.chunking.rounds_per_chunk}")
console.print(f" Index Mode: {self.config.index.mode}")
console.print(f" Max Iterations: {self.config.evaluation.max_iterations}")
if Confirm.ask("\nModify settings?"):
# Chunking settings
if Confirm.ask("Modify chunking settings?"):
rounds = Prompt.ask(
"Rounds per chunk",
default=str(self.config.chunking.rounds_per_chunk)
)
self.config.chunking.rounds_per_chunk = int(rounds)
overlap = Prompt.ask(
"Overlap rounds",
default=str(self.config.chunking.overlap_rounds)
)
self.config.chunking.overlap_rounds = int(overlap)
# Index settings
if Confirm.ask("Modify index settings?"):
mode = Prompt.ask(
"Index mode",
choices=["dense", "sparse", "hybrid"],
default=self.config.index.mode
)
self.config.index.mode = IndexMode(mode)
# Agent settings
if Confirm.ask("Modify agent settings?"):
max_iter = Prompt.ask(
"Max iterations",
default=str(self.config.evaluation.max_iterations)
)
self.config.evaluation.max_iterations = int(max_iter)
self.config.agent.enable_reasoning = Confirm.ask(
"Enable reasoning output?",
default=self.config.agent.enable_reasoning
)
# Save configuration
if Confirm.ask("Save configuration to file?"):
config_file = Prompt.ask("Config file path", default="config.json")
self.config.save(config_file)
console.print(f"[green]Configuration saved to {config_file}[/green]")
def evaluate_single_test(self):
"""Evaluate a single test case"""
if not self.test_cases_loaded:
console.print("[yellow]No test cases loaded. Please load test cases first.[/yellow]")
return
# Show available test cases for selection
console.print("\n[bold]Available Test Cases:[/bold]")
# Group by category for better organization
categories = {}
for test_id, test_case in self.evaluator.test_cases.items():
cat = test_case.category
if cat not in categories:
categories[cat] = []
categories[cat].append((test_id, test_case.title))
# Display test cases by category
test_ids_list = []
for cat in sorted(categories.keys()):
console.print(f"\n[cyan]{cat.upper()}:[/cyan]")
for test_id, title in sorted(categories[cat]):
test_ids_list.append(test_id)
# Show index number for easier selection
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]")
# Allow user to enter either test ID or number
user_input = Prompt.ask("Select test case")
# Check if user entered a number
test_id = None
if user_input.isdigit():
idx = int(user_input) - 1
if 0 <= idx < len(test_ids_list):
test_id = test_ids_list[idx]
else:
console.print(f"[red]Invalid selection number: {user_input}[/red]")
return
else:
# User entered a test ID directly
test_id = user_input
if test_id not in self.evaluator.test_cases:
console.print(f"[red]Test case {test_id} not found[/red]")
return
test_case = self.evaluator.test_cases[test_id]
console.print(f"\n[cyan]Evaluating: {test_case.title}[/cyan]")
console.print(f"Question: {test_case.user_question}\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(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("\nAgent Answer:")
console.print(Panel(result.agent_answer, border_style="cyan"))
console.print("\nEvaluation Criteria:")
console.print(Panel(result.evaluation_criteria, border_style="green"))
# Display LLM evaluation if available
if hasattr(result, 'llm_evaluation') and result.llm_evaluation and 'reward' in result.llm_evaluation:
console.print("\nLLM Evaluation:")
llm_eval = result.llm_evaluation
eval_color = "green" if llm_eval['passed'] else "red"
console.print(f" [{'green' if llm_eval['passed'] else 'red'}]Passed: {'Yes' if llm_eval['passed'] else 'No'}[/{eval_color}]")
console.print(f" Reward Score: {llm_eval['reward']:.3f}/1.000")
console.print(f" Reasoning: {llm_eval.get('reasoning', 'N/A')}")
if llm_eval.get('required_info_found'):
console.print("\n Required Information:")
for info, found in llm_eval['required_info_found'].items():
check = "[green]✓[/green]" if found else "[red]✗[/red]"
console.print(f" {check} {info}")
console.print(f"\nStatistics:")
console.print(f" Iterations: {result.iterations}")
console.print(f" Tool Calls: {result.tool_calls}")
console.print(f" Chunks Indexed: {result.chunk_count}")
console.print(f" Processing Time: {result.processing_time:.2f}s")
console.print(f" Indexing Time: {result.indexing_time:.2f}s")
except Exception as e:
progress.update(task, completed=True)
console.print(f"Error: {e}")
def evaluate_category(self):
"""Evaluate all test cases in a category"""
if not self.test_cases_loaded:
console.print("[yellow]No test cases loaded. Please load test cases first.[/yellow]")
return
# Show available categories and count
categories_count = {}
for test_case in self.evaluator.test_cases.values():
cat = test_case.category
categories_count[cat] = categories_count.get(cat, 0) + 1
console.print("\n[bold]Available Categories:[/bold]")
for cat in sorted(categories_count.keys()):
console.print(f" {cat}: {categories_count[cat]} test cases")
category = Prompt.ask(
"\nSelect category",
choices=list(sorted(categories_count.keys()))
)
# Get and display test cases in category
test_ids = []
console.print(f"\n[cyan]Test cases in {category}:[/cyan]")
for tid, tc in sorted(self.evaluator.test_cases.items()):
if tc.category == category:
test_ids.append(tid)
console.print(f" • {tid}: {tc.title[:60]}...")
console.print(f"\n[cyan]Total: {len(test_ids)} test cases[/cyan]")
if not Confirm.ask("Proceed with evaluation?"):
return
# Evaluate
with Progress(console=console) as progress:
task = progress.add_task(f"Evaluating {category}...", total=len(test_ids))
for test_id in test_ids:
try:
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)
console.print(f"[green]✓ Evaluation complete for {category}[/green]")
self.show_category_results(category)
def evaluate_all(self):
"""Evaluate all loaded test cases"""
if not self.test_cases_loaded:
console.print("[yellow]No test cases loaded. Please load test cases first.[/yellow]")
return
total = len(self.evaluator.test_cases)
console.print(f"\n[cyan]Will evaluate {total} test cases[/cyan]")
if not Confirm.ask("Proceed with full evaluation?"):
return
# Evaluate all
results = self.evaluator.evaluate_batch()
console.print(f"[green]✓ Evaluated {len(results)} test cases[/green]")
# Show summary
successful = sum(1 for r in results.values() if r.success)
console.print(f"Success rate: {successful}/{total} ({100*successful/total:.1f}%)")
def view_results(self):
"""View evaluation results"""
if not self.evaluator.results:
console.print("[yellow]No evaluation results available[/yellow]")
return
# Create results table
table = Table(title="Evaluation Results")
table.add_column("Test ID", style="cyan")
table.add_column("Category", style="magenta")
table.add_column("Status", justify="center")
table.add_column("Iterations", justify="right")
table.add_column("Tool Calls", justify="right")
table.add_column("Time (s)", justify="right")
for test_id, result in sorted(self.evaluator.results.items()):
test_case = self.evaluator.test_cases.get(test_id)
category = test_case.category if test_case else "unknown"
status = "[green]✓[/green]" if result.success else "[red]✗[/red]"
table.add_row(
test_id,
category,
status,
str(result.iterations),
str(result.tool_calls),
f"{result.processing_time:.2f}"
)
console.print(table)
# Summary statistics
total = len(self.evaluator.results)
successful = sum(1 for r in self.evaluator.results.values() if r.success)
console.print(f"\n[bold]Summary:[/bold]")
console.print(f" Total: {total}")
console.print(f" Successful: {successful} ({100*successful/total:.1f}%)")
# Option to view details
if Confirm.ask("\nView detailed result?"):
test_id = Prompt.ask("Enter test case ID")
if test_id in self.evaluator.results:
result = self.evaluator.results[test_id]
console.print(Panel(
f"[bold]Test: {test_id}[/bold]\n\n"
f"Status: {'Success' if result.success else 'Failed'}\n"
f"Iterations: {result.iterations}\n"
f"Tool Calls: {result.tool_calls}\n"
f"Chunks: {result.chunk_count}\n"
f"Processing Time: {result.processing_time:.2f}s\n"
f"Indexing Time: {result.indexing_time:.2f}s\n\n"
f"[yellow]Agent Answer:[/yellow]\n{result.agent_answer[:300]}...",
border_style="cyan"
))
def show_category_results(self, category: str):
"""Show results for a specific category"""
category_results = {
tid: r for tid, r in self.evaluator.results.items()
if tid in self.evaluator.test_cases and
self.evaluator.test_cases[tid].category == category
}
if not category_results:
console.print(f"[yellow]No results for {category}[/yellow]")
return
successful = sum(1 for r in category_results.values() if r.success)
total = len(category_results)
console.print(f"\n[bold]{category} Results:[/bold]")
console.print(f" Success Rate: {successful}/{total} ({100*successful/total:.1f}%)")
# Average metrics
if total > 0:
avg_iter = sum(r.iterations for r in category_results.values()) / total
avg_tools = sum(r.tool_calls for r in category_results.values()) / total
avg_time = sum(r.processing_time for r in category_results.values()) / total
console.print(f" Avg Iterations: {avg_iter:.1f}")
console.print(f" Avg Tool Calls: {avg_tools:.1f}")
console.print(f" Avg Time: {avg_time:.2f}s")
def generate_report(self):
"""Generate and save evaluation report"""
if not self.evaluator.results:
console.print("[yellow]No results to report[/yellow]")
return
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
report_file = f"results/report_{timestamp}.txt"
report = self.evaluator.generate_report(report_file)
console.print(f"[green]✓ Report saved to {report_file}[/green]")
if Confirm.ask("Display report?"):
console.print("\n" + report)
# Save JSON results
if Confirm.ask("Save detailed results to JSON?"):
json_file = f"results/results_{timestamp}.json"
self.evaluator.save_results(json_file)
console.print(f"[green]✓ Results saved to {json_file}[/green]")
def demo_mode(self):
"""Run a quick demo with a simple test case"""
console.print("\n[cyan]Demo Mode - Quick Test[/cyan]")
console.print("This will run a simple Layer 1 test case for demonstration.\n")
# Load only layer1 test cases
if not self.test_cases_loaded:
console.print("Loading Layer 1 test cases...")
test_cases = self.evaluator.load_test_cases("layer1")
self.test_cases_loaded = True
# Get first layer1 test case
layer1_tests = [
tid for tid, tc in self.evaluator.test_cases.items()
if tc.category == "layer1"
]
if not layer1_tests:
console.print("[red]No Layer 1 test cases available[/red]")
return
test_id = layer1_tests[0]
test_case = self.evaluator.test_cases[test_id]
console.print(f"[bold]Demo Test Case:[/bold] {test_case.title}")
console.print(f"[bold]Question:[/bold] {test_case.user_question}\n")
if Confirm.ask("Run demo?"):
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
console=console
) as progress:
task = progress.add_task("Running demo...", total=None)
try:
result = self.evaluator.evaluate_test_case(test_id)
progress.update(task, completed=True)
# Display results
console.print("\n[bold green]Demo Complete![/bold green]\n")
console.print(f"[bold]Agent's Answer:[/bold]")
console.print(Panel(result.agent_answer, border_style="cyan"))
console.print(f"\n[bold]Performance:[/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" Time: {result.processing_time:.2f}s")
except Exception as e:
progress.update(task, completed=True)
console.print(f"[red]Demo failed: {e}[/red]")
def _apply_cli_overrides(config: Config, args) -> Config:
"""把命令行参数覆盖到配置上（未指定的项保持默认，不改变原有行为）。"""
if args.provider:
config.llm.provider = args.provider
if args.model:
config.llm.model = args.model
if args.index_mode:
config.index.mode = IndexMode(args.index_mode)
if args.backend:
config.index.retrieval_backend = args.backend
if args.store_path:
config.index.index_path = args.store_path
if args.test_cases_dir:
config.evaluation.test_cases_dir = args.test_cases_dir
if args.rounds_per_chunk:
config.chunking.rounds_per_chunk = args.rounds_per_chunk
if args.top_k:
config.agent.max_search_results = args.top_k
return config
def main():
"""Main entry point"""
parser = argparse.ArgumentParser(
description="实验 3-10 · 智能体化 RAG 用户记忆评估系统",
epilog=(
"示例:\n"
" python main.py # 交互式菜单（默认）\n"
" python main.py --mode offline-demo # 离线对比演示，无需 API / port 4242\n"
" python main.py --mode batch --category layer2 --backend local\n"
" python main.py --mode batch --test-id layer2_01_multiple_vehicles --provider openai\n"
),
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--config", type=str,
help="配置文件（JSON）路径"
)
parser.add_argument(
"--mode",
choices=["interactive", "batch", "demo", "offline-demo"],
default="interactive",
help="运行模式：interactive 交互菜单（默认）/ batch 批量评估 / demo 快速演示 / offline-demo 离线检索对比"
)
parser.add_argument(
"--category",
choices=["layer1", "layer2", "layer3"],
help="批量模式下要评估的难度层次"
)
parser.add_argument(
"--test-id", type=str,
help="指定要评估的单个用例 ID"
)
parser.add_argument(
"--query", type=str,
help="offline-demo 模式下覆盖用例自带的用户问题"
)
parser.add_argument(
"--provider", type=str,
help="LLM 提供商（如 openai / kimi / siliconflow / deepseek 等）"
)
parser.add_argument(
"--model", type=str,
help="LLM 模型名，覆盖提供商默认模型"
)
parser.add_argument(
"--index-mode", choices=["dense", "sparse", "hybrid"],
help="检索策略：dense 稠密 / sparse 稀疏(BM25) / hybrid 混合"
)
parser.add_argument(
"--backend", choices=["auto", "local", "pipeline"],
help="检索后端：auto 自动（默认，pipeline 不可用则本地）/ local 内置离线 BM25 / pipeline 外部 4242 服务"
)
parser.add_argument(
"--top-k", type=int,
help="每次记忆检索返回的记忆块数量"
)
parser.add_argument(
"--rounds-per-chunk", type=int,
help="对话历史分块时每块的轮数（默认 20）"
)
parser.add_argument(
"--store-path", type=str,
help="记忆索引的存储路径前缀（默认 indexes/memory_index）"
)
parser.add_argument(
"--test-cases-dir", type=str,
help="评估集 test_cases 目录（默认 ../user-memory-evaluation/test_cases）"
)
parser.add_argument(
"--output", type=str,
help="结果输出文件路径"
)
args = parser.parse_args()
# offline-demo 模式：委托给完全离线的对比演示脚本
if args.mode == "offline-demo":
import offline_demo
demo_args = offline_demo.build_parser().parse_args([])
if args.test_id:
demo_args.test_id = args.test_id
if args.query:
demo_args.query = args.query
if args.top_k:
demo_args.top_k = args.top_k
if args.rounds_per_chunk:
demo_args.rounds_per_chunk = args.rounds_per_chunk
if args.test_cases_dir:
demo_args.test_cases_dir = args.test_cases_dir
if args.output:
demo_args.output = args.output
offline_demo.run_demo(demo_args)
return
# Load configuration
config = None
if args.config:
config = Config.load(args.config)
else:
config = Config.from_env()
# 应用命令行覆盖项
config = _apply_cli_overrides(config, args)
if args.mode == "interactive":
# Interactive mode
evaluator = InteractiveRAGEvaluator(config)
evaluator.run()
elif args.mode == "batch":
# Batch mode
evaluator = UserMemoryEvaluator(config)
if args.test_id:
# Evaluate single test
evaluator.load_test_cases()
result = evaluator.evaluate_test_case(args.test_id)
console.print(f"Result: {'Success' if result.success else 'Failed'}")
else:
# Evaluate category or all
evaluator.load_test_cases(args.category)
results = evaluator.evaluate_batch(category=args.category)
# Generate report
report_file = args.output or f"results/report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
report = evaluator.generate_report(report_file)
console.print(f"Report saved to {report_file}")
# Summary
total = len(results)
successful = sum(1 for r in results.values() if r.success)
console.print(f"Success rate: {successful}/{total} ({100*successful/total:.1f}%)")
elif args.mode == "demo":
# Demo mode
evaluator = InteractiveRAGEvaluator(config)
evaluator.demo_mode()
if __name__ == "__main__":
main()
offline_demo.py¶
#!/usr/bin/env python3
"""实验 3-10 离线演示：智能体化记忆检索 vs. 朴素单次检索
本脚本完全离线运行（不需要 port 4242 检索服务，也不需要任何 LLM API Key），
用来直观展示书中实验 3-10 的核心论点：
把用户的跨会话对话历史当作知识库、赋予 Agent「多轮迭代检索」能力后，
它能主动发现单次检索会遗漏的关键信息，从而在「第二层次·多会话检索」
任务上显著超过朴素的一次性检索（naive recall）。
演示载体是评估集里的 layer2_01_multiple_vehicles 用例：用户在两通不同电话里
分别聊到本田 Accord（已在 Firestone 预约周五保养）和特斯拉 Model 3（未预约）。
当用户问「我要给我的车约个保养，我都约了哪些服务？」时，正确回答必须同时覆盖
两辆车的服务状态——这正是「消歧」所要求的。
检索质量用「决定性证据召回率」度量：一次正确回答依赖若干条决定性事实
（如本田预约确认号 FS-447291、特斯拉作为第二辆车存在），我们统计每种策略
检索到的文本块是否覆盖了这些事实。所有数字都由真实的 BM25 检索计算得出，
不含任何人为编造。
"""
import argparse
import json
import logging
import re
import sys
from pathlib import Path
from typing import Dict, List, Optional
import yaml
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from config import Config
from chunker import ConversationChunker
from indexer import MemoryIndexer
console = Console()
# 默认用例及其决定性证据标记（每个标记代表回答必须覆盖的一条事实）。
# 标记只是「用哪个关键词定位承载该事实的文本块」，召回率由真实检索计算。
DEFAULT_TEST_ID = "layer2_01_multiple_vehicles"
DEFAULT_GOLD_MARKERS = {
"本田已确认预约(FS-447291)": "FS-447291",
"特斯拉Model 3作为第二辆车存在": "Model 3",
}
# 从检索结果里抽取「车辆实体」用的通用规则：年份+品牌+车型，或品牌+车型。
_ENTITY_PATTERN = re.compile(
r"\b((?:19|20)\d{2}\s+)?"
r"(Honda|Toyota|Tesla|Ford|BMW|Audi|Chevrolet|Nissan|Mazda|Subaru|Lexus|Kia|Hyundai)"
r"(?:\s+[A-Z][a-zA-Z0-9]+){0,2}"
)
def _load_test_case(test_cases_dir: Path, test_id: str) -> Optional[dict]:
"""在 test_cases_dir 下按 test_id 查找并加载 YAML 用例。"""
for yaml_file in test_cases_dir.rglob("*.yaml"):
try:
data = yaml.safe_load(yaml_file.read_text(encoding="utf-8"))
except Exception:
continue
if data and data.get("test_id") == test_id:
return data
return None
def _extract_vehicle_entities(texts: List[str], limit: int = 4) -> List[str]:
"""从一批文本里通用地抽取车辆实体短语（不依赖具体用例）。"""
seen: Dict[str, int] = {}
for text in texts:
for match in _ENTITY_PATTERN.finditer(text):
phrase = re.sub(r"\s+", " ", match.group(0)).strip()
# 去掉前导年份，聚焦「品牌 车型」，让二次查询更聚焦。
phrase = re.sub(r"^(?:19|20)\d{2}\s+", "", phrase)
seen[phrase] = seen.get(phrase, 0) + 1
# 按出现频次排序，取前 limit 个作为待追查的实体。
ranked = sorted(seen.items(), key=lambda kv: kv[1], reverse=True)
return [phrase for phrase, _ in ranked[:limit]]
def _covered_markers(retrieved_texts: List[str], gold_markers: Dict[str, str]) -> Dict[str, bool]:
"""判断每条决定性事实是否被检索到的文本覆盖。"""
joined = "\n".join(retrieved_texts)
return {name: (marker in joined) for name, marker in gold_markers.items()}
def naive_retrieval(indexer: MemoryIndexer, question: str, top_k: int) -> List:
"""基线：只用原始问题做一次检索。"""
return indexer.search(question, top_k=top_k)
def agentic_retrieval(indexer: MemoryIndexer, question: str, top_k: int,
max_followups: int, verbose: bool = True) -> Dict:
"""智能体化多轮检索（离线确定性模拟）。
模拟 Agent 的 ReAct 检索循环，但用规则代替 LLM，从而完全离线：
1. 用原始问题做首轮检索；
2. 从首轮结果里发现「还提到了哪些车辆实体」（关键线索）；
3. 针对每个发现的实体追加一次聚焦查询「<实体> service appointment scheduled」；
4. 合并所有轮次的结果作为最终检索集合。
"""
trace: List[Dict] = []
round1 = indexer.search(question, top_k=top_k)
retrieved = {r.chunk_id: r for r in round1}
trace.append({"query": question, "hits": [r.chunk_id for r in round1]})
if verbose:
console.print(f"[cyan] [第1轮] 查询:[/cyan] {question}")
console.print(f" 命中 {len(round1)} 块: "
+ ", ".join(f"{r.chunk.conversation_id}#{r.chunk.start_round}-{r.chunk.end_round}"
for r in round1))
# 从首轮结果里发现车辆实体（关键线索）
entities = _extract_vehicle_entities([r.chunk.to_text() for r in round1])
if verbose:
console.print(f"[cyan] [评估] 从首轮结果中发现车辆实体:[/cyan] {entities or '（无）'}")
for entity in entities[:max_followups]:
followup_query = f"{entity} service appointment scheduled"
hits = indexer.search(followup_query, top_k=top_k)
for r in hits:
retrieved.setdefault(r.chunk_id, r)
trace.append({"query": followup_query, "hits": [r.chunk_id for r in hits]})
if verbose:
console.print(f"[cyan] [追查] 查询:[/cyan] {followup_query}")
console.print(f" 命中 {len(hits)} 块: "
+ ", ".join(f"{r.chunk.conversation_id}#{r.chunk.start_round}-{r.chunk.end_round}"
for r in hits))
return {"results": list(retrieved.values()), "trace": trace}
def run_demo(args) -> Dict:
# 强制离线本地 BM25 后端
config = Config.from_env()
config.index.retrieval_backend = "local"
config.chunking.rounds_per_chunk = args.rounds_per_chunk
test_cases_dir = Path(args.test_cases_dir)
if not test_cases_dir.is_absolute():
test_cases_dir = (Path(__file__).parent / test_cases_dir).resolve()
console.print(Panel.fit(
"[bold cyan]实验 3-10 离线演示[/bold cyan]\n"
"智能体化记忆检索 vs. 朴素单次检索（本地 BM25，无需 API / port 4242）",
border_style="cyan"))
data = _load_test_case(test_cases_dir, args.test_id)
if not data:
console.print(f"[red]找不到用例 {args.test_id}，检索目录: {test_cases_dir}[/red]")
sys.exit(1)
question = args.query or data.get("user_question", "")
gold_markers = DEFAULT_GOLD_MARKERS if args.test_id == DEFAULT_TEST_ID else {}
if args.gold_marker:
gold_markers = {m: m for m in args.gold_marker}
# 分块 + 建索引（离线）
chunker = ConversationChunker(config.chunking)
chunks = chunker.chunk_test_case_conversations(data)
indexer = MemoryIndexer(config.index)
indexer.add_chunks(chunks)
# 多会话概览
sessions = {}
for c in chunks:
sessions.setdefault(c.conversation_id, []).append(c)
console.print(f"\n[bold]用例:[/bold] {data.get('title', args.test_id)}")
console.print(f"[bold]用户问题:[/bold] {question}")
console.print(f"[bold]跨会话记忆:[/bold] 共 {len(sessions)} 个历史会话，"
f"切分为 {len(chunks)} 个记忆块（每块 {args.rounds_per_chunk} 轮）")
for conv_id, cs in sessions.items():
meta = cs[0].metadata
console.print(f" • 会话 [magenta]{conv_id}[/magenta]（{meta.get('business', '?')} / "
f"{meta.get('department', '?')}）：{len(cs)} 块")
# 两种策略
console.print("\n[bold yellow]策略 A · 朴素单次检索（baseline）[/bold yellow]")
naive_results = naive_retrieval(indexer, question, args.top_k)
console.print(f" 单次查询命中 {len(naive_results)} 块: "
+ ", ".join(f"{r.chunk.conversation_id}#{r.chunk.start_round}-{r.chunk.end_round}"
for r in naive_results))
console.print("\n[bold green]策略 B · 智能体化多轮检索（memory-RAG）[/bold green]")
agentic = agentic_retrieval(indexer, question, args.top_k, args.max_followups)
agentic_results = agentic["results"]
# 计算决定性证据召回
naive_cover = _covered_markers([r.chunk.to_text() for r in naive_results], gold_markers)
agentic_cover = _covered_markers([r.chunk.to_text() for r in agentic_results], gold_markers)
def recall(cover: Dict[str, bool]) -> float:
return (sum(cover.values()) / len(cover)) if cover else 0.0
# 指标表
table = Table(title="检索质量对比（决定性证据召回）")
table.add_column("指标", style="cyan")
table.add_column("朴素单次检索", justify="center", style="yellow")
table.add_column("智能体化多轮检索", justify="center", style="green")
table.add_row("检索查询次数", "1", str(len(agentic["trace"])))
table.add_row("检索到的记忆块数", str(len(naive_results)), str(len(agentic_results)))
for name in gold_markers:
table.add_row(
f"覆盖事实: {name}",
"[green]✓[/green]" if naive_cover.get(name) else "[red]✗[/red]",
"[green]✓[/green]" if agentic_cover.get(name) else "[red]✗[/red]",
)
if gold_markers:
table.add_row("[bold]决定性证据召回率[/bold]",
f"[bold]{recall(naive_cover)*100:.0f}%[/bold]",
f"[bold]{recall(agentic_cover)*100:.0f}%[/bold]")
naive_ok = all(naive_cover.values())
agentic_ok = all(agentic_cover.values())
table.add_row("能否完整消歧作答",
"[green]能[/green]" if naive_ok else "[red]不能[/red]",
"[green]能[/green]" if agentic_ok else "[red]不能[/red]")
console.print()
console.print(table)
if gold_markers:
console.print(Panel(
"朴素单次检索只发一次查询，命中被「保养预约」这类关键词主导的文本块，"
"容易漏掉另一辆车的关键信息；智能体化检索从首轮结果里发现「还有第二辆车」，"
"再对每辆车追加聚焦查询，最终把两辆车的服务状态都取回，"
"从而在多会话检索任务上超过 naive recall。",
title="结论", border_style="green"))
result = {
"test_id": args.test_id,
"question": question,
"num_sessions": len(sessions),
"num_chunks": len(chunks),
"top_k": args.top_k,
"naive": {
"num_queries": 1,
"num_retrieved": len(naive_results),
"coverage": naive_cover,
"recall": recall(naive_cover),
},
"agentic": {
"num_queries": len(agentic["trace"]),
"num_retrieved": len(agentic_results),
"coverage": agentic_cover,
"recall": recall(agentic_cover),
"trace": agentic["trace"],
},
}
if args.output:
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
Path(args.output).write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
console.print(f"\n[green]✓ 结果已写入 {args.output}[/green]")
return result
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="实验 3-10 离线演示：智能体化记忆检索 vs. 朴素单次检索（本地 BM25，无需 API）",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument("--test-id", default=DEFAULT_TEST_ID,
help=f"评估集用例 ID（默认: {DEFAULT_TEST_ID}）")
parser.add_argument("--test-cases-dir", default="../user-memory-evaluation/test_cases",
help="评估集 test_cases 目录（默认: ../user-memory-evaluation/test_cases）")
parser.add_argument("--query", default=None,
help="覆盖用例自带的用户问题，指定要检索的问题")
parser.add_argument("--top-k", type=int, default=3,
help="每次检索返回的记忆块数量（默认: 3）")
parser.add_argument("--rounds-per-chunk", type=int, default=20,
help="对话历史分块时每块的轮数（默认: 20）")
parser.add_argument("--max-followups", type=int, default=4,
help="智能体化检索最多追加的聚焦查询数（默认: 4）")
parser.add_argument("--gold-marker", action="append", default=None,
help="自定义决定性证据关键词（可多次指定）；不指定时用内置默认")
parser.add_argument("--output", default=None,
help="将结构化结果写入指定 JSON 文件路径")
parser.add_argument("--quiet", action="store_true",
help="降低日志噪声（只显示 WARNING 及以上）")
return parser
def main():
args = build_parser().parse_args()
if args.quiet:
logging.getLogger().setLevel(logging.WARNING)
run_demo(args)
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_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_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_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_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")