contextual-retrieval¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/contextual-retrieval
项目说明¶
Contextual Retrieval System - Educational Implementation¶
An educational implementation of Anthropic's Contextual Retrieval technique, demonstrating how contextualizing chunks before indexing dramatically improves retrieval accuracy in RAG systems.
🌟 Key Insight¶
The Problem: Traditional RAG systems lose context when chunking documents. A chunk saying "The company's revenue grew by 3%" loses meaning without knowing which company or time period.
The Solution: Contextual Retrieval prepends chunk-specific explanatory context to each chunk before embedding and indexing, preserving semantic meaning.
📊 核心实验：离线量化召回提升（实验 3-11）¶
上面的 Key Insight 是本项目要验证的核心主张。compare_retrieval.py 用一组
可控的对比实验完全离线地量化它：把同一批文本块分别以两种方式建 BM25 索引——
无上下文（只索引原始文本）与有上下文（索引 LLM 生成的前缀 + 原始文本）——
再在评测集 evaluation/retrieval_eval.json（15 条查询 + 人工标注 gold 文本块）
上比较 recall@k（命中率）。无需任何 API 或检索服务（BM25 + jieba 分词）。
# 跑完整对比表（默认语料 document_store.json，默认评测集）
python compare_retrieval.py
# 查看每条查询的命中排名明细
python compare_retrieval.py --per-query
# 临时单条查询：并排看无上下文 / 有上下文的 Top-K 结果
python compare_retrieval.py --query "国家主席有哪些职权？" --top-k 5
# 只看无上下文基线 / 另存机器可读结果
python compare_retrieval.py --mode plain
python compare_retrieval.py --output result.json
# 中文 --help
python compare_retrieval.py --help
真实运行输出（22 个《宪法》《检察官法》文本块，15 条查询，jieba 分词）：
检索召回对比：无上下文分块 vs. 上下文感知检索（BM25）
====================================================================
方法 recall@1 recall@3 recall@5
----------------------------------------------------
无上下文 (plain) 60.0% 86.7% 93.3%
有上下文 (ctx) 86.7% 86.7% 93.3%
----------------------------------------------------
提升 (Δpp) +26.7pp +0.0pp +0.0pp
----------------------------------------------------
失败率下降 67% 0% 0%
结论与书中一致：为文本块补上上下文前缀显著提升 top-1 召回（60% → 86.7%，
失败率 1−recall@1 下降 67%），前缀为 BM25 注入了“身份标签”式的可匹配关键词。
这一提升在 recall@1 上最明显；--query 模式可直观看到前缀如何把“正确所属章节”
的文本块重新排到前面。
--method embedding/--method hybrid（上下文向量嵌入 + RRF 融合）需要 调用 embedding API，无法离线运行；脚本会给出提示并回退到 BM25 离线结果。 稠密检索 + 重排序的完整实现见contextual_tools.py。
同一对比逻辑也内建在 ContextualChunker.compare_retrieval_methods() 中（书中
compare_retrieval_methods 功能），可对任意一组 ContextualChunk 直接做并排检索对比。
📚 Educational Features¶
This implementation includes extensive logging and comparison capabilities to understand:
- How Context Generation Works: Watch the LLM generate context for each chunk
- Dual Indexing Strategy: See how both BM25 and embeddings benefit from context
- Comparison Mode: Run with
use_contextual=Falseto compare against standard chunking - Performance Metrics: Track improvements in retrieval accuracy
- Cost Analysis: Understand the token usage and costs
🏗️ Architecture¶
┌─────────────────────────────────────────┐
│ Document Input │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Basic Chunking │
│ (Respects paragraph boundaries) │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Context Generation (Optional) │
│ Using LLM API │
│ (Enabled with use_contextual=True) │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Enhanced Chunks │
│ • Contextual: Context + Original Text │
│ • Standard: Original Text Only │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Retrieval Pipeline Indexing │
│ • Sparse Index (BM25) │
│ • Dense Index (Embeddings) │
└────────────────┬────────────────────────┘
│
▼
┌─────────────────────────────────────────┐
│ Hybrid Search with Reranking │
│ Combines BM25 + Embedding scores │
│ Cross-encoder reranking for accuracy │
└─────────────────────────────────────────┘
🚀 Quick Start¶
1. Installation¶
2. Configure Environment¶
cp env.example .env
# Edit .env and add your API keys:
# - MOONSHOT_API_KEY for Kimi
# - ARK_API_KEY for Doubao
# - OPENAI_API_KEY for OpenAI
# - etc.
3. Start the Retrieval Pipeline¶
# In a separate terminal, start the retrieval pipeline server
cd ../retrieval-pipeline
python main.py
# Server will run on http://localhost:4242
4. Index Documents¶
# Index Chinese law documents with contextual enhancement
python index_local_laws_contextual.py
# Or index without contextual enhancement for comparison
python index_local_laws_contextual.py --no-contextual
5. Run Queries¶
# Interactive mode with contextual retrieval
python main.py
# Query with specific mode
python main.py --query "宪法第一条是什么" --mode agentic
# Compare agentic vs non-agentic modes
python main.py --query "宪法第一条是什么" --mode compare
Context Generation Process¶
The system generates context for each chunk by:
- Providing the full document (or surrounding context) to the LLM
- Showing the specific chunk to be contextualized
- Asking for concise context (2-3 sentences) that situates the chunk
Example prompt template:
<document>
[Full document or surrounding context]
</document>
Here is the chunk we want to situate:
<chunk>
[Specific chunk text]
</chunk>
Please give a short, succinct context to situate this chunk within the overall document...
📚 References¶
🤝 Contributing¶
This is an educational implementation. Contributions welcome for: - Additional chunking strategies - Alternative context generation prompts - Performance optimizations - Evaluation metrics - Visualization tools
📝 License¶
Educational project for learning purposes.
🙏 Acknowledgments¶
Based on research by Anthropic's engineering team on improving RAG retrieval accuracy through contextual enhancement.
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 System with 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, LLMConfig, AgentConfig
from tools import KnowledgeBaseTools, get_tool_definitions
def _reasoning_safe_temperature(model, requested=1.0):
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
Return 1 for those; otherwise the requested value so non-reasoning
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
m = str(model or "").lower().replace("/", "-")
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
@dataclass
class Message:
"""Represents a message in the conversation"""
role: str # "user", "assistant", "tool"
content: str
tool_calls: Optional[List[Dict[str, Any]]] = None
tool_call_id: Optional[str] = None
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
class AgenticRAG:
"""Agentic RAG system with ReAct pattern and multiple LLM provider support"""
def __init__(self, config: Optional[Config] = None):
"""Initialize the agent"""
self.config = config or Config.from_env()
# Initialize LLM client
self._init_llm_client()
# Initialize knowledge base tools
self.kb_tools = KnowledgeBaseTools(self.config.knowledge_base)
# Conversation history
self.conversation_history: List[Dict[str, Any]] = []
# Tool definitions
self.tools = get_tool_definitions()
logger.info(f"Initialized AgenticRAG 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) -> str:
"""Generate the system prompt"""
return """You are an intelligent assistant with access to a knowledge base. Your primary role is to answer questions accurately based on the information available in the knowledge base.
## Important Guidelines:
1. **Knowledge Base Only**: You MUST only answer questions based on information found in the knowledge base. If the information is not available, clearly state that you cannot answer based on the available knowledge.
2. **Use Tools Effectively**:
- Use `knowledge_base_search` to search for relevant information
- Use `get_document` to retrieve complete documents when you need more context
- You may need multiple searches with different queries to fully answer a question
3. **Citations Required**: Always include citations in your answers. Format citations as [Doc: document_id] or [Chunk: chunk_id] inline with your response.
4. **Reasoning Process**: Think step-by-step:
- First, understand what information is needed
- Search for relevant information
- If needed, retrieve full documents for context
- Synthesize the information to answer the question
- Include proper citations
5. **Handle Follow-ups**: For follow-up questions, consider the conversation context but always verify information from the knowledge base.
6. **Be Accurate**: Never make up information. If something is unclear or not found, say so explicitly.
Remember: Your credibility depends on providing accurate, well-cited information from the knowledge base only."""
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""Execute a tool and return the result"""
try:
if tool_name == "knowledge_base_search":
query = arguments.get("query", "")
if self.config.agent.verbose:
logger.info(f"Executing tool: {tool_name} with args: {arguments}")
results = self.kb_tools.knowledge_base_search(query)
if not results:
logger.info(f"No results found for query: {query}")
return {"status": "no_results", "message": f"No relevant documents found for query: {query}"}
# Format results for agent
formatted_results = []
for r in results:
formatted_results.append({
"doc_id": r["doc_id"],
"chunk_id": r["chunk_id"],
"text": r["text"],
"score": r["score"]
})
return {
"status": "success",
"results": formatted_results,
"total_found": len(results)
}
elif tool_name == "get_document":
doc_id = arguments.get("doc_id", "")
document = self.kb_tools.get_document(doc_id)
if "error" in document:
return {"status": "error", "message": document["error"]}
return {
"status": "success",
"document": {
"doc_id": document.get("doc_id", doc_id),
"content": document.get("content", ""),
"metadata": document.get("metadata", {})
}
}
else:
return {"status": "error", "message": f"Unknown tool: {tool_name}"}
except Exception as e:
logger.error(f"Tool execution error: {e}")
return {"status": "error", "message": str(e)}
def _build_messages(self, user_query: str) -> List[Dict[str, Any]]:
"""Build messages for the LLM including conversation history"""
messages = [{"role": "system", "content": self._get_system_prompt()}]
# Add conversation history (limited)
history_limit = self.config.agent.conversation_history_limit
if len(self.conversation_history) > history_limit:
messages.extend(self.conversation_history[-history_limit:])
else:
messages.extend(self.conversation_history)
# Add current user query
messages.append({"role": "user", "content": user_query})
return messages
def query(self, user_query: str, stream: bool = None) -> Any:
"""
Process a user query using the ReAct pattern.
Args:
user_query: The user's question
stream: Whether to stream the response
Returns:
The agent's response (string or generator for streaming)
"""
if stream is None:
stream = self.config.llm.stream
# Build messages
messages = self._build_messages(user_query)
# Track iterations
iterations = 0
max_iterations = self.config.agent.max_iterations
# Process with ReAct loop
while iterations < max_iterations:
iterations += 1
if self.config.agent.verbose:
logger.info(f"Iteration {iterations}/{max_iterations}")
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 # We handle streaming separately
)
message = response.choices[0].message
# 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
]
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:
logger.error(f"Failed to parse tool arguments: {tool_call.function.arguments}")
arguments = {}
# Execute tool
result = self._execute_tool(tool_name, arguments)
if self.config.agent.verbose:
logger.info(f"Tool result: {json.dumps(result, indent=2, ensure_ascii=False)}")
# Add tool result to messages
tool_message = {
"role": "tool",
"tool_call_id": tool_call.id,
"content": json.dumps(result, indent=2, ensure_ascii=False)
}
messages.append(tool_message)
# Continue loop for next iteration
continue
else:
# No tool calls, we have final answer
# Update conversation history
self.conversation_history.append({"role": "user", "content": user_query})
self.conversation_history.append(assistant_msg)
final_response = message.content or ""
# Log the final response if verbose
if self.config.agent.verbose:
logger.info(f"Final response generated (length: {len(final_response)} chars)")
# Return response
if stream:
return self._stream_response(final_response)
else:
return final_response
except Exception as e:
logger.error(f"Error in query processing: {e}")
error_msg = f"Error processing query: {str(e)}"
if stream:
return self._stream_response(error_msg)
else:
return error_msg
# Max iterations reached
logger.warning(f"Max iterations ({max_iterations}) reached")
final_msg = "I need more iterations to fully answer your question. Please try rephrasing or breaking down your query."
if stream:
return self._stream_response(final_msg)
else:
return final_msg
def _stream_response(self, content: str) -> Generator[str, None, None]:
"""Stream response content"""
# Simple character streaming for demonstration
for char in content:
yield char
def query_non_agentic(self, user_query: str, stream: bool = None) -> Any:
"""
Non-agentic RAG mode: Simple retrieval + LLM response.
Args:
user_query: The user's question
stream: Whether to stream the response
Returns:
The response (string or generator for streaming)
"""
if stream is None:
stream = self.config.llm.stream
try:
# Simple retrieval
if self.config.agent.verbose:
logger.info(f"Non-agentic mode: searching for '{user_query}'")
search_results = self.kb_tools.knowledge_base_search(user_query)
if self.config.agent.verbose:
logger.info(f"Non-agentic mode: found {len(search_results)} results")
# Build context from search results
context_parts = []
for i, result in enumerate(search_results[:5], 1): # Top 5 results
context_parts.append(
f"[Document {i}] (ID: {result['doc_id']}, Chunk: {result['chunk_id']})\n{result['text']}\n"
)
if not context_parts:
context = "No relevant information found in the knowledge base."
else:
context = "\n".join(context_parts)
# Build prompt
system_prompt = """You are an assistant that answers questions based on provided context from a knowledge base.
IMPORTANT RULES:
1. Only answer based on the provided context
2. Include citations in format [Doc: document_id]
3. If the context doesn't contain the answer, say so clearly
4. Be accurate and don't make up information"""
user_prompt = f"""Context from knowledge base:
{context}
User Question: {user_query}
Please answer the question based only on the provided context. Include citations."""
# Call LLM
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
if stream:
response_stream = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
max_tokens=self.config.llm.max_tokens,
stream=True
)
def response_generator():
for chunk in response_stream:
if chunk.choices[0].delta.content:
yield chunk.choices[0].delta.content
return response_generator()
else:
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
max_tokens=self.config.llm.max_tokens,
stream=False
)
return response.choices[0].message.content
except Exception as e:
logger.error(f"Error in non-agentic query: {e}")
error_msg = f"Error processing query: {str(e)}"
if stream:
return self._stream_response(error_msg)
else:
return error_msg
def clear_history(self):
"""Clear conversation history"""
self.conversation_history = []
logger.info("Conversation history cleared")
chunking.py¶
"""Document chunking and indexing script"""
import os
import json
import hashlib
import logging
import requests
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
from datetime import datetime
from config import ChunkingConfig, KnowledgeBaseConfig, KnowledgeBaseType
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
class DocumentChunker:
"""Document chunking with configurable strategies"""
def __init__(self, config: Optional[ChunkingConfig] = None):
self.config = config or ChunkingConfig()
def chunk_text(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""
Chunk text into smaller segments.
Args:
text: Document text to chunk
doc_id: Document identifier
Returns:
List of chunks with metadata
"""
chunks = []
if self.config.respect_paragraph_boundary:
chunks = self._chunk_by_paragraphs(text, doc_id)
else:
chunks = self._chunk_by_size(text, doc_id)
logger.info(f"Created {len(chunks)} chunks for document {doc_id}")
return chunks
def _chunk_by_paragraphs(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""Chunk text respecting paragraph boundaries"""
paragraphs = text.split('\n\n')
chunks = []
current_chunk = []
current_size = 0
for para in paragraphs:
para = para.strip()
if not para:
continue
para_size = len(para)
# If single paragraph exceeds max size, split it
if para_size > self.config.max_chunk_size:
# Save current chunk if exists
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
current_chunk = []
current_size = 0
# Split large paragraph
sentences = self._split_into_sentences(para)
for sent in sentences:
if len(sent) > self.config.max_chunk_size:
# Force split very long sentences
for i in range(0, len(sent), self.config.chunk_size):
sub_chunk = sent[i:i + self.config.chunk_size]
chunks.append(self._create_chunk(sub_chunk, doc_id, len(chunks)))
else:
chunks.append(self._create_chunk(sent, doc_id, len(chunks)))
continue
# Check if adding this paragraph exceeds chunk size
if current_size + para_size > self.config.chunk_size and current_chunk:
# Save current chunk
chunk_text = '\n\n'.join(current_chunk)
chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
# Start new chunk with overlap
if self.config.chunk_overlap > 0 and current_chunk:
# Keep last paragraph for overlap
current_chunk = [current_chunk[-1], para]
current_size = len(current_chunk[0]) + para_size
else:
current_chunk = [para]
current_size = para_size
else:
current_chunk.append(para)
current_size += para_size
# Save final chunk
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
if len(chunk_text) >= self.config.min_chunk_size:
chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
return chunks
def _chunk_by_size(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""Simple size-based chunking"""
chunks = []
for i in range(0, len(text), self.config.chunk_size - self.config.chunk_overlap):
chunk_text = text[i:i + self.config.chunk_size]
if len(chunk_text) >= self.config.min_chunk_size:
chunks.append(self._create_chunk(chunk_text, doc_id, len(chunks)))
return chunks
def _split_into_sentences(self, text: str) -> List[str]:
"""Split text into sentences (simple implementation)"""
# Simple sentence splitting for Chinese and English
import re
# Split on common sentence endings
sentences = re.split(r'([。！？\.!?]+)', text)
# Reconstruct sentences with their endings
result = []
for i in range(0, len(sentences) - 1, 2):
if i + 1 < len(sentences):
result.append(sentences[i] + sentences[i + 1])
else:
result.append(sentences[i])
return [s.strip() for s in result if s.strip()]
def _create_chunk(self, text: str, doc_id: str, chunk_index: int) -> Dict[str, Any]:
"""Create a chunk with metadata"""
chunk_id = f"{doc_id}_chunk_{chunk_index}"
return {
"chunk_id": chunk_id,
"doc_id": doc_id,
"text": text,
"chunk_index": chunk_index,
"char_count": len(text),
"hash": hashlib.md5(text.encode()).hexdigest()
}
class DocumentIndexer:
"""Index documents to knowledge base"""
def __init__(self,
kb_config: Optional[KnowledgeBaseConfig] = None,
chunking_config: Optional[ChunkingConfig] = None):
self.kb_config = kb_config or KnowledgeBaseConfig()
self.chunker = DocumentChunker(chunking_config)
self.indexed_docs = {}
def index_file(self, file_path: str, doc_id: Optional[str] = None) -> Dict[str, Any]:
"""
Index a single file.
Args:
file_path: Path to the file
doc_id: Optional document ID
Returns:
Indexing result
"""
file_path = Path(file_path)
if not file_path.exists():
return {"error": f"File not found: {file_path}"}
# Generate doc_id if not provided
if not doc_id:
doc_id = file_path.stem
# Read file content
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
except Exception as e:
return {"error": f"Error reading file: {e}"}
# Chunk the document
chunks = self.chunker.chunk_text(content, doc_id)
# Index chunks
result = self._index_chunks(chunks, doc_id, content)
# Store full document
self._store_document(doc_id, content, {"source_file": str(file_path)})
return result
def index_directory(self, dir_path: str, extensions: List[str] = None) -> Dict[str, Any]:
"""
Index all files in a directory.
Args:
dir_path: Directory path
extensions: File extensions to include (e.g., ['.txt', '.md'])
Returns:
Indexing results
"""
dir_path = Path(dir_path)
if not dir_path.exists():
return {"error": f"Directory not found: {dir_path}"}
extensions = extensions or ['.txt', '.md', '.json']
results = {"indexed": [], "errors": []}
for file_path in dir_path.rglob('*'):
if file_path.is_file() and file_path.suffix in extensions:
doc_id = f"{file_path.parent.name}/{file_path.stem}"
result = self.index_file(str(file_path), doc_id)
if "error" in result:
results["errors"].append({
"file": str(file_path),
"error": result["error"]
})
else:
results["indexed"].append({
"file": str(file_path),
"doc_id": doc_id,
"chunks": result.get("chunks_indexed", 0)
})
logger.info(f"Indexed {len(results['indexed'])} files, {len(results['errors'])} errors")
return results
def _index_chunks(self, chunks: List[Dict[str, Any]], doc_id: str, full_content: str) -> Dict[str, Any]:
"""Index chunks to the knowledge base"""
if self.kb_config.type == KnowledgeBaseType.LOCAL:
return self._index_to_local(chunks, doc_id)
elif self.kb_config.type == KnowledgeBaseType.DIFY:
return self._index_to_dify(chunks, doc_id, full_content)
else:
return {"error": f"Unsupported KB type: {self.kb_config.type}"}
def _index_to_local(self, chunks: List[Dict[str, Any]], doc_id: str) -> Dict[str, Any]:
"""Index to local retrieval pipeline"""
indexed_count = 0
errors = []
for chunk in chunks:
try:
# Index each chunk
response = requests.post(
f"{self.kb_config.local_base_url}/index",
json={
"text": chunk["text"],
"doc_id": chunk["doc_id"],
"metadata": {
"chunk_id": chunk["chunk_id"],
"chunk_index": chunk["chunk_index"],
"char_count": chunk["char_count"]
}
}
)
response.raise_for_status()
indexed_count += 1
except Exception as e:
errors.append(f"Error indexing chunk {chunk['chunk_id']}: {e}")
result = {
"doc_id": doc_id,
"chunks_indexed": indexed_count,
"total_chunks": len(chunks)
}
if errors:
result["errors"] = errors
return result
def _index_to_dify(self, chunks: List[Dict[str, Any]], doc_id: str, full_content: str) -> Dict[str, Any]:
"""Index to Dify knowledge base"""
if not self.kb_config.dify_api_key:
return {"error": "Dify API key not configured"}
try:
headers = {
"Authorization": f"Bearer {self.kb_config.dify_api_key}",
"Content-Type": "application/json"
}
# Dify expects documents, not individual chunks
# So we'll create segments from our chunks
segments = []
for chunk in chunks:
segments.append({
"content": chunk["text"],
"keywords": [], # Can add keywords if needed
"enabled": True
})
payload = {
"name": doc_id,
"text": full_content,
"indexing_technique": "high_quality", # or "economy"
"process_rule": {
"mode": "custom",
"rules": {
"pre_processing_rules": [],
"segmentation": {
"separator": "\n\n",
"max_tokens": self.chunker.config.chunk_size // 4 # Rough token estimate
}
}
}
}
if self.kb_config.dify_dataset_id:
# Add to existing dataset
response = requests.post(
f"{self.kb_config.dify_base_url}/datasets/{self.kb_config.dify_dataset_id}/documents",
headers=headers,
json=payload
)
else:
# Create new document
response = requests.post(
f"{self.kb_config.dify_base_url}/documents",
headers=headers,
json=payload
)
response.raise_for_status()
return {
"doc_id": doc_id,
"chunks_indexed": len(chunks),
"total_chunks": len(chunks),
"dify_response": response.json()
}
except Exception as e:
return {"error": f"Error indexing to Dify: {e}"}
def _store_document(self, doc_id: str, content: str, metadata: Dict[str, Any]):
"""Store full document locally"""
# Store in local file for retrieval
store_path = self.kb_config.document_store_path
try:
# Load existing store
if os.path.exists(store_path):
with open(store_path, 'r', encoding='utf-8') as f:
store = json.load(f)
else:
store = {}
# Add document
store[doc_id] = {
"doc_id": doc_id,
"content": content,
"metadata": metadata,
"indexed_at": datetime.now().isoformat()
}
# Save store
with open(store_path, 'w', encoding='utf-8') as f:
json.dump(store, f, ensure_ascii=False, indent=2)
self.indexed_docs[doc_id] = True
logger.info(f"Stored document {doc_id}")
except Exception as e:
logger.error(f"Error storing document: {e}")
def main():
"""Main function for standalone chunking and indexing"""
import argparse
from config import Config
parser = argparse.ArgumentParser(description="Chunk and index documents")
parser.add_argument("path", help="File or directory path to index")
parser.add_argument("--chunk-size", type=int, default=2048, help="Chunk size in characters")
parser.add_argument("--max-chunk-size", type=int, default=1024, help="Max chunk size")
parser.add_argument("--overlap", type=int, default=200, help="Chunk overlap")
parser.add_argument("--kb-type", choices=["local", "dify"], default="local", help="Knowledge base type")
parser.add_argument("--extensions", nargs="+", default=[".txt", ".md"], help="File extensions to index")
args = parser.parse_args()
# Create config
config = Config.from_env()
config.chunking.chunk_size = args.chunk_size
config.chunking.max_chunk_size = args.max_chunk_size
config.chunking.chunk_overlap = args.overlap
config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
# Create indexer
indexer = DocumentIndexer(config.knowledge_base, config.chunking)
# Index path
path = Path(args.path)
if path.is_file():
result = indexer.index_file(str(path))
elif path.is_dir():
result = indexer.index_directory(str(path), args.extensions)
else:
print(f"Path not found: {path}")
return
# Print results
print(json.dumps(result, indent=2, ensure_ascii=False))
if __name__ == "__main__":
main()
compare_retrieval.py¶
#!/usr/bin/env python3
"""上下文感知检索对比评测（实验 3-11）
本脚本用可控的对比实验量化“上下文感知检索”相较传统分块的检索召回提升：
同一批文本块分别以两种方式建立 BM25 索引——
* 无上下文（plain） ：只索引原始文本块 metadata.original_text
* 有上下文（contextual）：索引 LLM 生成的前缀 + 原始文本块（content 字段）
然后在同一评测集上比较 recall@k（命中率：前 k 个结果中是否含有相关文本块）。
这正是 Anthropic “Contextual Retrieval” 的核心主张：为文本块补上上下文前缀，
能同时增强 BM25（稀疏）与向量（稠密）检索的召回率。
BM25 检索完全离线，无需任何 API 或检索服务；embedding / hybrid 方法需要
调用 embedding API（见 --method 说明）。
用法示例：
python compare_retrieval.py # 用默认评测集跑对比表
python compare_retrieval.py --query "国家主席有哪些职权？" # 单条查询并排对比
python compare_retrieval.py --mode plain # 只看无上下文基线
python compare_retrieval.py --output result.json # 另存机器可读结果
"""
import argparse
import json
import sys
from pathlib import Path
from typing import Dict, List, Optional
import numpy as np
from rank_bm25 import BM25Okapi
try:
import jieba
if hasattr(jieba, "setLogLevel"):
jieba.setLogLevel(60) # 关闭 jieba 的加载日志
_HAS_JIEBA = True
except Exception: # pragma: no cover - jieba 一般随 requirements 安装
_HAS_JIEBA = False
# ---------------------------------------------------------------------------
# 分词：中文没有空格，直接 .split() 会把整段当成一个 token，BM25 完全失效。
# 默认用 jieba 分词；--no-jieba 时退化为字符二元组（bigram），同样可离线运行。
# ---------------------------------------------------------------------------
def tokenize(text: str, use_jieba: bool = True) -> List[str]:
"""把文本切成 token 列表，供 BM25 使用。"""
text = (text or "").lower()
if use_jieba and _HAS_JIEBA:
return [t for t in jieba.cut(text) if t.strip()]
# 退化方案：中文字符二元组 + 连续 ASCII 词
tokens: List[str] = []
buf = ""
chars = list(text)
for ch in chars:
if ch.isascii() and (ch.isalnum()):
buf += ch
continue
if buf:
tokens.append(buf)
buf = ""
if not ch.isspace():
tokens.append(ch)
if buf:
tokens.append(buf)
# 追加中文 bigram，提升匹配粒度
cjk = [c for c in text if "一" <= c <= "鿿"]
tokens.extend(cjk[i] + cjk[i + 1] for i in range(len(cjk) - 1))
return tokens
# ---------------------------------------------------------------------------
# 语料加载
# ---------------------------------------------------------------------------
def load_corpus(path: str) -> List[Dict]:
"""从 document_store.json 载入分块，返回 [{chunk_id, contextual, plain, context}]。
每个分块的 content 字段是“上下文前缀 + 原始文本”，metadata.original_text
是不带上下文的原始文本，正好用于两种索引方式的对照。
"""
with open(path, "r", encoding="utf-8") as f:
store = json.load(f)
chunks: List[Dict] = []
for chunk_id, entry in store.items():
if "_chunk_" not in chunk_id:
continue # 跳过整篇文档条目
if not isinstance(entry, dict):
continue
meta = entry.get("metadata", {}) or {}
contextual_text = entry.get("content", "") or ""
plain_text = meta.get("original_text") or contextual_text
# 上下文前缀 = contextual 去掉结尾的 original_text
context = contextual_text
if plain_text and contextual_text.endswith(plain_text):
context = contextual_text[: len(contextual_text) - len(plain_text)].strip()
chunks.append({
"chunk_id": chunk_id,
"contextual": contextual_text,
"plain": plain_text,
"context": context,
})
return chunks
def load_eval(path: str) -> List[Dict]:
"""载入评测集，返回 [{id, query, gold_chunk_id, ...}]。"""
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
return data.get("queries", data if isinstance(data, list) else [])
# ---------------------------------------------------------------------------
# BM25 检索器
# ---------------------------------------------------------------------------
class BM25Retriever:
"""对给定文本字段建立 BM25 索引的简单检索器。"""
def __init__(self, chunks: List[Dict], field: str, use_jieba: bool = True):
self.chunk_ids = [c["chunk_id"] for c in chunks]
self.use_jieba = use_jieba
corpus_tokens = [tokenize(c[field], use_jieba) for c in chunks]
self.index = BM25Okapi(corpus_tokens)
def rank(self, query: str) -> List[str]:
"""返回按相关性从高到低排序的 chunk_id 列表。"""
scores = self.index.get_scores(tokenize(query, self.use_jieba))
order = np.argsort(scores)[::-1]
return [self.chunk_ids[i] for i in order]
def scored(self, query: str, top_k: int) -> List[Dict]:
"""返回前 top_k 个结果及其分数。"""
scores = self.index.get_scores(tokenize(query, self.use_jieba))
order = np.argsort(scores)[::-1][:top_k]
return [{"chunk_id": self.chunk_ids[i], "score": float(scores[i])} for i in order]
# ---------------------------------------------------------------------------
# 评测
# ---------------------------------------------------------------------------
def recall_at_k(retriever: BM25Retriever, queries: List[Dict], ks: List[int]) -> Dict:
"""计算一批查询在各 k 值下的 recall@k（命中率）。"""
per_query = []
hits = {k: 0 for k in ks}
for q in queries:
ranking = retriever.rank(q["query"])
gold = q["gold_chunk_id"]
rank_pos = ranking.index(gold) + 1 if gold in ranking else None
row = {"id": q.get("id"), "query": q["query"], "gold": gold, "rank": rank_pos}
for k in ks:
hit = rank_pos is not None and rank_pos <= k
row[f"hit@{k}"] = hit
if hit:
hits[k] += 1
per_query.append(row)
n = len(queries)
recall = {k: (hits[k] / n if n else 0.0) for k in ks}
return {"recall": recall, "per_query": per_query, "n": n}
def print_comparison_table(plain: Optional[Dict], contextual: Optional[Dict], ks: List[int]):
"""打印 recall@k 对比表。"""
print("\n" + "=" * 68)
print("检索召回对比：无上下文分块 vs. 上下文感知检索（BM25）")
print("=" * 68)
header = " k | " + " | ".join(f"{'无上下文':>10}" if False else f"recall@{k:<3}" for k in ks)
# 逐行打印每个方法
col_w = 12
line = f"{'方法':<16}" + "".join(f"recall@{k}".rjust(col_w) for k in ks)
print(line)
print("-" * len(line))
if plain:
print(f"{'无上下文 (plain)':<16}" + "".join(f"{plain['recall'][k]*100:>10.1f}%" for k in ks))
if contextual:
print(f"{'有上下文 (ctx)':<16}" + "".join(f"{contextual['recall'][k]*100:>10.1f}%" for k in ks))
if plain and contextual:
print("-" * len(line))
deltas = []
for k in ks:
d = (contextual["recall"][k] - plain["recall"][k]) * 100
deltas.append(f"{d:>+9.1f}pp")
print(f"{'提升 (Δpp)':<16}" + "".join(s.rjust(col_w) for s in deltas))
# 检索失败率下降（对应书中“1 - recall@k”口径）
print("-" * len(line))
fails = []
for k in ks:
p_fail = 1 - plain["recall"][k]
c_fail = 1 - contextual["recall"][k]
if p_fail > 0:
red = (p_fail - c_fail) / p_fail * 100
fails.append(f"{red:>9.0f}%")
else:
fails.append(f"{'-':>10}")
print(f"{'失败率下降':<16}" + "".join(s.rjust(col_w) for s in fails))
print("=" * 68)
def print_per_query(result: Dict, label: str):
print(f"\n[{label}] 每条查询命中排名（rank=gold 文本块在结果中的名次，— 表示未召回）")
for row in result["per_query"]:
print(f" {row['id']} rank={str(row['rank']):>3} gold={row['gold']:<28} {row['query'][:32]}")
# ---------------------------------------------------------------------------
# 单条查询并排对比
# ---------------------------------------------------------------------------
def single_query_compare(chunks: List[Dict], query: str, top_k: int, use_jieba: bool,
mode: str):
id2chunk = {c["chunk_id"]: c for c in chunks}
def show(field_label, field):
retr = BM25Retriever(chunks, field, use_jieba)
print(f"\n[{field_label}] Top-{top_k}")
print("-" * 60)
for i, r in enumerate(retr.scored(query, top_k), 1):
c = id2chunk[r["chunk_id"]]
snippet = c["plain"].replace("<!-- FORCE BREAK -->", "").replace("\n", " ").strip()[:48]
ctx = c["context"].replace("\n", " ").strip()[:40]
print(f" {i}. score={r['score']:6.2f} {r['chunk_id']}")
if field == "contextual" and ctx:
print(f" 上下文前缀: {ctx}")
print(f" 原文: {snippet}")
print("\n" + "=" * 60)
print(f"查询: {query}")
print("=" * 60)
if mode in ("plain", "both"):
show("无上下文 (plain)", "plain")
if mode in ("contextual", "both"):
show("有上下文 (contextual)", "contextual")
# ---------------------------------------------------------------------------
# 可选：embedding / hybrid（需要 API）
# ---------------------------------------------------------------------------
def embedding_unavailable_notice(method: str):
print(f"\n[提示] --method {method} 需要调用 embedding API（稠密向量），无法离线运行。")
print(" 请在 .env 中配置 OPENAI_API_KEY / SILICONFLOW_API_KEY 等，")
print(" 并使用 contextual_tools.ContextualKnowledgeBaseTools 的 embedding/hybrid 检索。")
print(" 本脚本的默认 --method bm25 已可完整复现书中“上下文增强 BM25”的召回提升结论。")
def build_arg_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="上下文感知检索对比评测：量化上下文前缀对检索召回（recall@k）的提升（实验 3-11）",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="示例:\n"
" python compare_retrieval.py\n"
" python compare_retrieval.py --query \"国家主席有哪些职权？\" --top-k 5\n"
" python compare_retrieval.py --mode both --k 1 3 5 --output result.json",
)
p.add_argument("--corpus", default="document_store.json",
help="语料文件（含 content 与 metadata.original_text 的分块存储），默认 document_store.json")
p.add_argument("--eval", dest="eval_path", default="evaluation/retrieval_eval.json",
help="评测集（query + gold_chunk_id），默认 evaluation/retrieval_eval.json")
p.add_argument("--query", default=None,
help="临时单条查询：并排展示无上下文/有上下文的 Top-K 检索结果（不跑整个评测集）")
p.add_argument("--mode", choices=["plain", "contextual", "both"], default="both",
help="对比哪种索引：plain=仅无上下文，contextual=仅有上下文，both=两者对比（默认）")
p.add_argument("--method", choices=["bm25", "embedding", "hybrid"], default="bm25",
help="检索方法：bm25（离线，默认）；embedding/hybrid 需 embedding API")
p.add_argument("--k", nargs="+", type=int, default=[1, 3, 5],
help="评测的 k 值列表（recall@k），默认 1 3 5")
p.add_argument("--top-k", type=int, default=5,
help="--query 单查询模式下每种方法展示的结果条数，默认 5")
p.add_argument("--model", default=None,
help="embedding 模型名（仅 --method embedding/hybrid 时生效）")
p.add_argument("--no-jieba", action="store_true",
help="禁用 jieba 分词，改用字符二元组分词（无需 jieba 依赖）")
p.add_argument("--output", default=None,
help="将机器可读的评测结果写入该 JSON 文件")
p.add_argument("--per-query", action="store_true",
help="额外打印每条查询的命中排名明细")
return p
def main():
args = build_arg_parser().parse_args()
use_jieba = not args.no_jieba
corpus_path = Path(args.corpus)
if not corpus_path.exists():
print(f"[错误] 找不到语料文件: {corpus_path}", file=sys.stderr)
sys.exit(1)
chunks = load_corpus(str(corpus_path))
if not chunks:
print(f"[错误] 语料中没有可用分块（缺少 *_chunk_* 条目）: {corpus_path}", file=sys.stderr)
sys.exit(1)
print(f"已加载 {len(chunks)} 个文本块 | 分词: {'jieba' if (use_jieba and _HAS_JIEBA) else '字符bigram'} "
f"| 检索方法: {args.method}")
if args.method in ("embedding", "hybrid"):
embedding_unavailable_notice(args.method)
# 仍继续用 BM25 给出可运行的离线结果
print(" 以下改用 BM25 给出离线对照结果。\n")
# 单条查询模式
if args.query:
single_query_compare(chunks, args.query, args.top_k, use_jieba, args.mode)
return
# 评测集模式
eval_path = Path(args.eval_path)
if not eval_path.exists():
print(f"[错误] 找不到评测集: {eval_path}", file=sys.stderr)
sys.exit(1)
queries = load_eval(str(eval_path))
ks = sorted(set(args.k))
plain_res = contextual_res = None
if args.mode in ("plain", "both"):
plain_res = recall_at_k(BM25Retriever(chunks, "plain", use_jieba), queries, ks)
if args.mode in ("contextual", "both"):
contextual_res = recall_at_k(BM25Retriever(chunks, "contextual", use_jieba), queries, ks)
print(f"评测集: {eval_path} 共 {len(queries)} 条查询")
print_comparison_table(plain_res, contextual_res, ks)
if args.per_query:
if plain_res:
print_per_query(plain_res, "无上下文 plain")
if contextual_res:
print_per_query(contextual_res, "有上下文 contextual")
if args.output:
out = {
"corpus": str(corpus_path),
"eval": str(eval_path),
"num_chunks": len(chunks),
"num_queries": len(queries),
"tokenizer": "jieba" if (use_jieba and _HAS_JIEBA) else "char-bigram",
"k": ks,
"plain": plain_res,
"contextual": contextual_res,
}
with open(args.output, "w", encoding="utf-8") as f:
json.dump(out, f, ensure_ascii=False, indent=2)
print(f"\n结果已写入 {args.output}")
if __name__ == "__main__":
main()
config.py¶
"""Configuration for Agentic RAG System"""
import os
from dataclasses import dataclass, field
from typing import Optional, Dict, Any
from enum import Enum
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 KnowledgeBaseType(str, Enum):
"""Knowledge base backend types"""
LOCAL = "local" # Local retrieval pipeline
DIFY = "dify" # Dify knowledge base API
RAPTOR = "raptor" # RAPTOR tree-based index
GRAPHRAG = "graphrag" # GraphRAG graph-based index
@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 = 1024
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"
}
}
@classmethod
def get_api_key(cls, provider: str) -> Optional[str]:
"""Get API key from environment"""
env_mappings = {
"siliconflow": "SILICONFLOW_API_KEY",
"doubao": "ARK_API_KEY",
"kimi": "MOONSHOT_API_KEY",
"moonshot": "MOONSHOT_API_KEY",
"openrouter": "OPENROUTER_API_KEY",
"openai": "OPENAI_API_KEY",
"groq": "GROQ_API_KEY",
"together": "TOGETHER_API_KEY",
"deepseek": "DEEPSEEK_API_KEY"
}
return os.getenv(env_mappings.get(provider.lower(), ""))
def get_client_config(self) -> Dict[str, Any]:
"""Get OpenAI client configuration"""
provider_lower = self.provider.lower()
defaults = self.PROVIDER_DEFAULTS.get(provider_lower, {})
# Get API key
api_key = self.api_key or self.get_api_key(provider_lower)
# Universal OpenRouter fallback: primary provider key absent but
# OPENROUTER_API_KEY present -> route through OpenRouter. Additionally,
# gpt-5.x (incl. gpt-5.6*) needs OpenAI org-verification on the direct
# API, so prefer OpenRouter for those ids whenever an OR key is present.
model_name = self.model or defaults.get("model")
openrouter_key = os.getenv("OPENROUTER_API_KEY")
prefer_openrouter = bool(openrouter_key) and str(model_name or "").lower().startswith("gpt-5")
if (not api_key or prefer_openrouter) and provider_lower != "openrouter" and openrouter_key:
model = _openrouter_model_id(model_name)
return {
"api_key": openrouter_key,
"base_url": "https://openrouter.ai/api/v1",
}, model
if not api_key:
raise ValueError(
f"API key required for provider '{provider_lower}'. Set the "
f"provider's key (e.g. MOONSHOT_API_KEY / OPENAI_API_KEY) or "
f"OPENROUTER_API_KEY to use the OpenRouter fallback."
)
# Build config
config = {
"api_key": api_key,
"model": self.model or defaults.get("model")
}
# Add base_url if not OpenAI
if "base_url" in defaults:
config["base_url"] = defaults["base_url"]
return config, config.pop("model")
@dataclass
class KnowledgeBaseConfig:
"""Knowledge base configuration"""
type: KnowledgeBaseType = KnowledgeBaseType.LOCAL
# Local retrieval pipeline config
local_base_url: str = "http://localhost:4242"
local_top_k: int = 3
# Dify config
dify_api_key: Optional[str] = field(default_factory=lambda: os.getenv("DIFY_API_KEY"))
dify_base_url: str = "https://api.dify.ai/v1"
dify_dataset_id: Optional[str] = None
dify_top_k: int = 10
# RAPTOR tree-based index config
raptor_base_url: str = "http://localhost:4242"
raptor_top_k: int = 10
raptor_search_levels: bool = True # Search across multiple tree levels
# GraphRAG graph-based index config
graphrag_base_url: str = "http://localhost:4242"
graphrag_top_k: int = 10
graphrag_search_type: str = "hybrid" # entity, community, or hybrid
# Document storage
document_store_path: str = "document_store.json"
@dataclass
class ChunkingConfig:
"""Document chunking configuration"""
chunk_size: int = 2048 # Characters per chunk
max_chunk_size: int = 1024 # Max size when respecting paragraph boundaries
chunk_overlap: int = 200 # Overlap between chunks
respect_paragraph_boundary: bool = True
min_chunk_size: int = 100 # Minimum chunk size
@dataclass
class AgentConfig:
"""Agent configuration"""
max_iterations: int = 10 # Max reasoning iterations
enable_reasoning_trace: bool = True
enable_citations: bool = True
strict_knowledge_base: bool = True # Only answer from knowledge base
conversation_history_limit: int = 20 # Max conversation turns to keep
verbose: bool = True
@dataclass
class EvaluationConfig:
"""Evaluation configuration"""
dataset_path: str = "evaluation/legal_qa_dataset.json"
results_path: str = "evaluation/results"
metrics: list = field(default_factory=lambda: ["accuracy", "relevance", "citation_quality"])
@dataclass
class Config:
"""Main configuration"""
llm: LLMConfig = field(default_factory=LLMConfig)
knowledge_base: KnowledgeBaseConfig = field(default_factory=KnowledgeBaseConfig)
chunking: ChunkingConfig = field(default_factory=ChunkingConfig)
agent: AgentConfig = field(default_factory=AgentConfig)
evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
@classmethod
def from_env(cls) -> "Config":
"""Create config from environment variables"""
config = cls()
# Override from env
if provider := os.getenv("LLM_PROVIDER"):
config.llm.provider = provider
if model := os.getenv("LLM_MODEL"):
config.llm.model = model
if kb_type := os.getenv("KB_TYPE"):
config.knowledge_base.type = KnowledgeBaseType(kb_type.lower())
return config
contextual_agent.py¶
"""Contextual Agentic RAG System
This module extends the base AgenticRAG to use contextual retrieval,
demonstrating the improved answer quality from better retrieval.
"""
import json
import logging
from typing import List, Dict, Any, Optional, Generator
from datetime import datetime
from agent import AgenticRAG, Message
from config import Config
from contextual_tools import ContextualKnowledgeBaseTools, ContextualSearchResult
logger = logging.getLogger(__name__)
class ContextualAgenticRAG(AgenticRAG):
"""
Enhanced Agentic RAG with contextual retrieval support.
Educational Features:
- Shows how better retrieval leads to better answers
- Logs retrieval quality metrics
- Compares contextual vs non-contextual results
"""
def __init__(self,
config: Optional[Config] = None,
kb_tools: Optional[ContextualKnowledgeBaseTools] = None,
use_contextual: bool = True):
"""
Initialize contextual agent.
Args:
config: Configuration object
kb_tools: Contextual knowledge base tools
use_contextual: Whether to use contextual retrieval
"""
self.use_contextual = use_contextual
# Initialize base class but skip KB tools initialization
self.config = config or Config.from_env()
self._init_llm_client()
# Use provided contextual KB tools or create new ones
if kb_tools:
self.kb_tools = kb_tools
else:
self.kb_tools = ContextualKnowledgeBaseTools(
self.config.knowledge_base,
use_contextual=use_contextual,
enable_comparison=False
)
# Initialize other components
self.conversation_history = []
self.tools = self._get_contextual_tool_definitions()
# Track retrieval metrics
self.retrieval_metrics = {
"queries_executed": 0,
"total_chunks_retrieved": 0,
"avg_retrieval_score": 0.0,
"contextual_chunks_used": 0,
"non_contextual_chunks_used": 0
}
mode = "CONTEXTUAL" if use_contextual else "NON-CONTEXTUAL"
logger.info(f"Initialized ContextualAgenticRAG in {mode} mode")
def _get_system_prompt(self) -> str:
"""Enhanced system prompt that leverages contextual information"""
base_prompt = super()._get_system_prompt()
if self.use_contextual:
contextual_addition = """
## Using Contextual Information:
When you receive search results, pay attention to the contextual information provided. Each chunk may include:
- Context that situates the chunk within the larger document
- Information about what section or topic the chunk belongs to
- Related entities and concepts mentioned
Use this contextual information to:
1. Better understand the relevance of each chunk
2. Identify connections between different pieces of information
3. Provide more accurate and complete answers
Remember that contextual chunks have been enhanced with additional information to improve retrieval accuracy."""
return base_prompt + contextual_addition
return base_prompt
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
"""
Execute tool with enhanced logging for educational purposes.
This override adds detailed logging to show how contextual
retrieval improves the search results.
"""
logger.debug(f"Executing tool: {tool_name} with args: {arguments}")
try:
if tool_name == "contextual_knowledge_search":
query = arguments.get("query", "")
method = arguments.get("method", "hybrid")
top_k = arguments.get("top_k", 20)
logger.info(f"\n{'='*60}")
logger.info(f"CONTEXTUAL SEARCH EXECUTION")
logger.info(f"{'='*60}")
logger.info(f"Query: {query}")
logger.info(f"Method: {method}")
logger.info(f"Top-K: {top_k}")
logger.info(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
# Perform search
results = self.kb_tools.contextual_search(query, method, top_k)
if not results:
logger.info("No results found")
return {"status": "no_results", "message": "No relevant documents found"}
# Log retrieval quality
avg_score = sum(r.score for r in results) / len(results)
logger.info(f"\nRetrieval Statistics:")
logger.info(f" Results found: {len(results)}")
logger.info(f" Average score: {avg_score:.4f}")
logger.info(f" Top score: {results[0].score:.4f}")
# Show distribution of retrieval methods if hybrid
if method == "hybrid":
bm25_only = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score == 0)
embedding_only = sum(1 for r in results if r.embedding_score > 0 and r.bm25_score == 0)
both = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score > 0)
logger.info(f"\nRetrieval Method Distribution:")
logger.info(f" BM25 only: {bm25_only}")
logger.info(f" Embedding only: {embedding_only}")
logger.info(f" Both methods: {both}")
# Update metrics
self.retrieval_metrics["queries_executed"] += 1
self.retrieval_metrics["total_chunks_retrieved"] += len(results)
# Calculate running average score
prev_avg = self.retrieval_metrics["avg_retrieval_score"]
n = self.retrieval_metrics["queries_executed"]
self.retrieval_metrics["avg_retrieval_score"] = (
(prev_avg * (n - 1) + avg_score) / n
)
if self.use_contextual:
self.retrieval_metrics["contextual_chunks_used"] += len(results)
else:
self.retrieval_metrics["non_contextual_chunks_used"] += len(results)
# Format results for agent
formatted_results = []
for i, r in enumerate(results, 1):
result_dict = {
"rank": i,
"doc_id": r.doc_id,
"chunk_id": r.chunk_id,
"text": r.text,
"score": r.score
}
formatted_results.append(result_dict)
# Log top results
if i <= 3:
logger.info(f"\nTop Result {i}:")
logger.info(f" Score: {r.score:.4f}")
logger.info(f" Text: {r.text}")
logger.info(f"{'='*60}\n")
return {
"status": "success",
"results": formatted_results,
"total_found": len(results),
"avg_score": avg_score,
"search_method": method,
"is_contextual": self.use_contextual
}
elif tool_name == "get_document":
# Use parent implementation
return super()._execute_tool(tool_name, arguments)
else:
return {"status": "error", "message": f"Unknown tool: {tool_name}"}
except Exception as e:
logger.error(f"Tool execution error: {e}")
return {"status": "error", "message": str(e)}
def _get_contextual_tool_definitions(self) -> List[Dict[str, Any]]:
"""Get tool definitions with contextual search capabilities"""
return [
{
"type": "function",
"function": {
"name": "contextual_knowledge_search",
"description": "Search the knowledge base using contextual or non-contextual retrieval. Returns chunks with optional context information that helps understand their relevance.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Natural language search query"
},
"method": {
"type": "string",
"enum": ["bm25", "embedding", "hybrid"],
"description": "Search method to use (default: hybrid)",
"default": "hybrid"
},
"top_k": {
"type": "integer",
"description": "Number of results to retrieve (default: 20)",
"default": 20
}
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "get_document",
"description": "Retrieve the complete content of a specific document",
"parameters": {
"type": "object",
"properties": {
"doc_id": {
"type": "string",
"description": "Document ID to retrieve"
}
},
"required": ["doc_id"]
}
}
}
]
def query(self, user_query: str, stream: bool = None) -> Any:
"""
Process query with contextual retrieval and detailed logging.
This override adds educational logging to show the retrieval process.
"""
logger.info(f"\n{'='*80}")
logger.info(f"PROCESSING QUERY: {user_query}")
logger.info(f"Mode: {'CONTEXTUAL' if self.use_contextual else 'NON-CONTEXTUAL'}")
logger.info(f"{'='*80}\n")
# Call parent implementation
result = super().query(user_query, stream)
# Log final metrics
logger.info(f"\n{'='*60}")
logger.info("QUERY PROCESSING COMPLETE")
logger.info(f"{'='*60}")
logger.info("Retrieval Metrics:")
logger.info(f" Queries executed: {self.retrieval_metrics['queries_executed']}")
logger.info(f" Total chunks retrieved: {self.retrieval_metrics['total_chunks_retrieved']}")
logger.info(f" Average retrieval score: {self.retrieval_metrics['avg_retrieval_score']:.4f}")
if self.use_contextual:
logger.info(f" Contextual chunks used: {self.retrieval_metrics['contextual_chunks_used']}")
else:
logger.info(f" Non-contextual chunks used: {self.retrieval_metrics['non_contextual_chunks_used']}")
logger.info(f"{'='*60}\n")
return result
def get_retrieval_metrics(self) -> Dict[str, Any]:
"""Get detailed retrieval metrics for analysis"""
metrics = self.retrieval_metrics.copy()
# Add KB statistics
kb_stats = self.kb_tools.get_statistics()
metrics["kb_stats"] = kb_stats
# Add mode information
metrics["mode"] = "contextual" if self.use_contextual else "non_contextual"
metrics["timestamp"] = datetime.now().isoformat()
return metrics
def compare_modes(self, user_query: str) -> Dict[str, Any]:
"""
Compare contextual vs non-contextual modes for the same query.
Educational method to demonstrate the difference in retrieval quality.
"""
comparison = {
"query": user_query,
"timestamp": datetime.now().isoformat(),
"modes": {}
}
# Test contextual mode
logger.info("Testing CONTEXTUAL mode...")
self.use_contextual = True
self.kb_tools.use_contextual = True
contextual_response = self.query(user_query, stream=False)
comparison["modes"]["contextual"] = {
"response": contextual_response,
"metrics": self.get_retrieval_metrics()
}
# Reset metrics
self.retrieval_metrics = {
"queries_executed": 0,
"total_chunks_retrieved": 0,
"avg_retrieval_score": 0.0,
"contextual_chunks_used": 0,
"non_contextual_chunks_used": 0
}
# Test non-contextual mode
logger.info("Testing NON-CONTEXTUAL mode...")
self.use_contextual = False
self.kb_tools.use_contextual = False
non_contextual_response = self.query(user_query, stream=False)
comparison["modes"]["non_contextual"] = {
"response": non_contextual_response,
"metrics": self.get_retrieval_metrics()
}
# Analyze differences
contextual_score = comparison["modes"]["contextual"]["metrics"]["avg_retrieval_score"]
non_contextual_score = comparison["modes"]["non_contextual"]["metrics"]["avg_retrieval_score"]
improvement = ((contextual_score - non_contextual_score) / non_contextual_score * 100) if non_contextual_score > 0 else 0
comparison["analysis"] = {
"score_improvement_pct": improvement,
"contextual_avg_score": contextual_score,
"non_contextual_avg_score": non_contextual_score,
"recommendation": (
"Use contextual retrieval" if improvement > 10
else "Consider cost-benefit of contextual retrieval"
)
}
return comparison
contextual_chunking.py¶
"""Contextual Chunking Module - Educational implementation of Anthropic's Contextual Retrieval
This module demonstrates the key insight from Anthropic's research:
- Traditional RAG destroys context by chunking documents
- Contextual Retrieval prepends chunk-specific context before embedding
- This preserves semantic meaning that would otherwise be lost
"""
import json
import hashlib
import logging
import requests
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass, field
import time
from openai import OpenAI
from config import ChunkingConfig, KnowledgeBaseConfig, KnowledgeBaseType, LLMConfig
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 ContextualChunk:
"""Enhanced chunk with contextual information"""
chunk_id: str
doc_id: str
text: str # Original chunk text
context: str # Generated contextual description
contextualized_text: str # Context + original text
chunk_index: int
char_count: int
context_tokens: int = 0 # Track token usage
generation_time: float = 0.0 # Track generation time
metadata: Dict[str, Any] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
return {
"chunk_id": self.chunk_id,
"doc_id": self.doc_id,
"text": self.text,
"context": self.context,
"contextualized_text": self.contextualized_text,
"chunk_index": self.chunk_index,
"char_count": self.char_count,
"context_tokens": self.context_tokens,
"generation_time": self.generation_time,
"metadata": self.metadata
}
class ContextualChunker:
"""
Implements contextual chunking inspired by Anthropic's Contextual Retrieval.
Key Educational Points:
1. Context Generation: Uses LLM to generate chunk-specific context
2. Prepending Strategy: Context is prepended to chunks before embedding
3. BM25 Enhancement: Contextual chunks improve both semantic and lexical search
4. Cost Optimization: Uses caching strategies to reduce API costs
"""
def __init__(self,
chunking_config: Optional[ChunkingConfig] = None,
llm_config: Optional[LLMConfig] = None,
use_contextual: bool = True):
"""
Initialize the contextual chunker.
Args:
chunking_config: Configuration for chunking parameters
llm_config: LLM configuration for context generation
use_contextual: Whether to generate contextual chunks (for comparison)
"""
self.chunking_config = chunking_config or ChunkingConfig()
self.llm_config = llm_config or LLMConfig()
self.use_contextual = use_contextual
# Initialize LLM client for context generation
if self.use_contextual:
self._init_llm_client()
# Statistics tracking
self.stats = {
"total_chunks": 0,
"contextual_chunks": 0,
"total_context_tokens": 0,
"total_generation_time": 0.0,
"cache_hits": 0,
"cache_misses": 0
}
# Context cache to avoid regenerating for similar chunks
self.context_cache = {}
logger.info(f"Initialized ContextualChunker (contextual={use_contextual})")
def _init_llm_client(self):
"""Initialize LLM client for context generation"""
client_config, model = self.llm_config.get_client_config()
base_url = client_config.pop("base_url", None)
if base_url:
self.client = OpenAI(base_url=base_url, **client_config)
else:
self.client = OpenAI(**client_config)
self.model = model
logger.info(f"Using {self.llm_config.provider} ({self.model}) for context generation")
def chunk_document(self,
text: str,
doc_id: str,
doc_metadata: Optional[Dict[str, Any]] = None,
on_chunk_ready: Optional[callable] = None) -> List[ContextualChunk]:
"""
Chunk a document with optional contextual enhancement.
Educational Note:
This is the core innovation - each chunk gets contextualized
with information about its position and meaning within the document.
Args:
text: Full document text
doc_id: Document identifier
doc_metadata: Optional document metadata
Returns:
List of ContextualChunk objects
"""
logger.info(f"Starting chunking for document {doc_id}")
start_time = time.time()
# Step 1: Create basic chunks
basic_chunks = self._create_basic_chunks(text, doc_id)
logger.info(f"Created {len(basic_chunks)} basic chunks")
# Step 2: Generate contextual enhancements if enabled
if self.use_contextual:
contextual_chunks = self._generate_contextual_chunks(
basic_chunks, text, doc_id, doc_metadata, on_chunk_ready
)
else:
# Create non-contextual chunks for comparison
contextual_chunks = []
for chunk in basic_chunks:
contextual_chunks.append(ContextualChunk(
chunk_id=chunk["chunk_id"],
doc_id=chunk["doc_id"],
text=chunk["text"],
context="", # No context in non-contextual mode
contextualized_text=chunk["text"], # Just the original text
chunk_index=chunk["chunk_index"],
char_count=chunk["char_count"],
metadata={"contextual": False}
))
# Update statistics
self.stats["total_chunks"] += len(contextual_chunks)
if self.use_contextual:
self.stats["contextual_chunks"] += len(contextual_chunks)
elapsed = time.time() - start_time
logger.info(f"Chunking completed in {elapsed:.2f}s")
logger.info(f"Statistics: {json.dumps(self.stats, indent=2)}")
return contextual_chunks
def _create_basic_chunks(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""Create basic chunks using traditional chunking"""
chunks = []
if self.chunking_config.respect_paragraph_boundary:
chunks = self._chunk_by_paragraphs(text, doc_id)
else:
chunks = self._chunk_by_size(text, doc_id)
return chunks
def _chunk_by_paragraphs(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""Chunk text respecting paragraph boundaries"""
paragraphs = text.split('\n\n')
chunks = []
current_chunk = []
current_size = 0
for para in paragraphs:
para = para.strip()
if not para:
continue
para_size = len(para)
# Handle oversized paragraphs
if para_size > self.chunking_config.max_chunk_size:
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
current_chunk = []
current_size = 0
# Split large paragraph into sentences
sentences = self._split_into_sentences(para)
for sent in sentences:
if len(sent) > self.chunking_config.max_chunk_size:
# Force split very long sentences
for i in range(0, len(sent), self.chunking_config.chunk_size):
sub_chunk = sent[i:i + self.chunking_config.chunk_size]
chunks.append(self._create_basic_chunk(sub_chunk, doc_id, len(chunks)))
else:
chunks.append(self._create_basic_chunk(sent, doc_id, len(chunks)))
continue
# Check if adding this paragraph exceeds chunk size
if current_size + para_size > self.chunking_config.chunk_size and current_chunk:
chunk_text = '\n\n'.join(current_chunk)
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
# Start new chunk with overlap
if self.chunking_config.chunk_overlap > 0 and current_chunk:
current_chunk = [current_chunk[-1], para]
current_size = len(current_chunk[0]) + para_size
else:
current_chunk = [para]
current_size = para_size
else:
current_chunk.append(para)
current_size += para_size
# Save final chunk
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
if len(chunk_text) >= self.chunking_config.min_chunk_size:
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
return chunks
def _chunk_by_size(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
"""Simple size-based chunking"""
chunks = []
for i in range(0, len(text), self.chunking_config.chunk_size - self.chunking_config.chunk_overlap):
chunk_text = text[i:i + self.chunking_config.chunk_size]
if len(chunk_text) >= self.chunking_config.min_chunk_size:
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
return chunks
def _split_into_sentences(self, text: str) -> List[str]:
"""Split text into sentences"""
import re
# Handle both English and Chinese sentence endings
sentences = re.split(r'([。！？\.!?]+)', text)
# Reconstruct sentences with their endings
result = []
for i in range(0, len(sentences) - 1, 2):
if i + 1 < len(sentences):
result.append(sentences[i] + sentences[i + 1])
else:
result.append(sentences[i])
return [s.strip() for s in result if s.strip()]
def _create_basic_chunk(self, text: str, doc_id: str, chunk_index: int) -> Dict[str, Any]:
"""Create a basic chunk dictionary"""
chunk_id = f"{doc_id}_chunk_{chunk_index}"
return {
"chunk_id": chunk_id,
"doc_id": doc_id,
"text": text,
"chunk_index": chunk_index,
"char_count": len(text),
"hash": hashlib.md5(text.encode()).hexdigest()
}
def _generate_contextual_chunks(self,
basic_chunks: List[Dict[str, Any]],
full_document: str,
doc_id: str,
doc_metadata: Optional[Dict[str, Any]] = None,
on_chunk_ready: Optional[callable] = None) -> List[ContextualChunk]:
"""
Generate contextual enhancements for chunks.
Educational Note:
Following Anthropic's Contextual Retrieval approach:
- Each chunk gets a concise context explaining its position in the document
- Context is prepended to the chunk before embedding
- This dramatically improves retrieval accuracy
"""
contextual_chunks = []
# No document summary needed - Anthropic's approach doesn't use it
doc_summary = None
for i, chunk in enumerate(basic_chunks):
logger.info(f"Generating context for chunk {i+1}/{len(basic_chunks)}")
# Check cache first
chunk_hash = chunk["hash"]
if chunk_hash in self.context_cache:
context = self.context_cache[chunk_hash]
generation_time = 0.0
context_tokens = 0
self.stats["cache_hits"] += 1
logger.debug(f"Cache hit for chunk {chunk['chunk_id']}")
else:
# Generate new context using Anthropic's approach
context, context_tokens, generation_time = self._generate_chunk_context(
chunk["text"],
full_document
)
# Cache the context
self.context_cache[chunk_hash] = context
self.stats["cache_misses"] += 1
self.stats["total_context_tokens"] += context_tokens
self.stats["total_generation_time"] += generation_time
# Create contextual chunk
contextualized_text = f"{context}\n\n{chunk['text']}" if context else chunk["text"]
contextual_chunk = ContextualChunk(
chunk_id=chunk["chunk_id"],
doc_id=doc_id,
text=chunk["text"],
context=context,
contextualized_text=contextualized_text,
chunk_index=chunk["chunk_index"],
char_count=len(contextualized_text),
context_tokens=context_tokens,
generation_time=generation_time,
metadata={
"contextual": True,
"original_char_count": chunk["char_count"],
"context_char_count": len(context)
}
)
contextual_chunks.append(contextual_chunk)
# Call the callback immediately if provided
if on_chunk_ready:
try:
on_chunk_ready(contextual_chunk)
except Exception as e:
logger.error(f"Error in on_chunk_ready callback: {e}")
# Log progress
if (i + 1) % 10 == 0:
avg_time = self.stats["total_generation_time"] / (i + 1)
logger.info(f"Progress: {i+1}/{len(basic_chunks)} chunks, avg time: {avg_time:.2f}s")
return contextual_chunks
def _generate_document_summary(self, document: str, doc_id: str) -> str:
"""
DEPRECATED: Not used in Anthropic's Contextual Retrieval approach.
Educational Note:
Anthropic's research shows that document summaries don't significantly
improve retrieval. Instead, they provide the full document directly
when generating chunk-specific context. This allows the LLM to understand
the exact context needed for each specific chunk.
"""
# This method is kept for backward compatibility but returns empty string
return ""
def _generate_chunk_context(self,
chunk_text: str,
full_document: str) -> Tuple[str, int, float]:
"""
Generate contextual description for a chunk using Anthropic's exact approach.
This follows Anthropic's Contextual Retrieval template exactly:
1. Provide the full document
2. Show the specific chunk
3. Ask for concise context to situate the chunk
Returns:
Tuple of (context, token_count, generation_time)
"""
start_time = time.time()
try:
# Use the exact prompt from Anthropic's blog post
# with added instruction to use the same language as the document
prompt = f"""<document>
{full_document}
</document>
Here is the chunk we want to situate within the whole document
<chunk>
{chunk_text}
</chunk>
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else. You MUST use the same language as the document."""
# Use the exact approach from Anthropic - no system message needed
response = self.client.chat.completions.create(
model=self.model,
messages=[
{"role": "user", "content": prompt}
],
temperature=_reasoning_safe_temperature(self.model, 0.3), # Low temperature for consistency
max_tokens=100 # Anthropic mentions 50-100 tokens typically
)
context = response.choices[0].message.content.strip()
# Estimate token count (rough approximation)
token_count = len(prompt.split()) + len(context.split())
generation_time = time.time() - start_time
logger.info(f"Generated context in {generation_time:.2f}s: {context}")
return context, token_count, generation_time
except Exception as e:
logger.error(f"Error generating chunk context: {e}")
return "", 0, time.time() - start_time
def compare_retrieval_methods(self,
query: str,
contextual_chunks: List[ContextualChunk],
non_contextual_chunks: List[ContextualChunk],
top_k: int = 5) -> Dict[str, Any]:
"""
Compare contextual vs non-contextual retrieval on the SAME query.
Educational Note:
This is the ``compare_retrieval_methods`` capability referenced in
实验 3-11. It builds two BM25 indexes fully offline (no API / server):
* contextual index over ``contextualized_text`` (前缀 + 原文)
* plain index over the original chunk ``text``
and returns the top-k ranked chunks under each, so the caller can see
exactly how the contextual prefix re-ranks the same corpus.
"""
from rank_bm25 import BM25Okapi
import numpy as np
from compare_retrieval import tokenize
results = {
"query": query,
"timestamp": datetime.now().isoformat(),
"contextual_results": [],
"non_contextual_results": [],
"analysis": {}
}
def _rank(chunks: List[ContextualChunk], field: str):
if not chunks:
return []
corpus = [tokenize(getattr(c, field)) for c in chunks]
index = BM25Okapi(corpus)
scores = index.get_scores(tokenize(query))
order = np.argsort(scores)[::-1][:top_k]
ranked = []
for rank, idx in enumerate(order, 1):
c = chunks[idx]
ranked.append({
"chunk_id": c.chunk_id,
"score": float(scores[idx]),
"rank": rank,
"text": c.text[:200],
"context": c.context[:200],
})
return ranked
results["contextual_results"] = _rank(contextual_chunks, "contextualized_text")
results["non_contextual_results"] = _rank(non_contextual_chunks, "text")
ctx_top = results["contextual_results"][0] if results["contextual_results"] else None
plain_top = results["non_contextual_results"][0] if results["non_contextual_results"] else None
results["analysis"] = {
"contextual_top_chunk": ctx_top["chunk_id"] if ctx_top else None,
"non_contextual_top_chunk": plain_top["chunk_id"] if plain_top else None,
"contextual_top_score": ctx_top["score"] if ctx_top else 0.0,
"non_contextual_top_score": plain_top["score"] if plain_top else 0.0,
"top1_changed": bool(ctx_top and plain_top and ctx_top["chunk_id"] != plain_top["chunk_id"]),
}
logger.info(f"Compared retrieval methods for query: {query} | "
f"top1 changed={results['analysis']['top1_changed']}")
return results
def get_statistics(self) -> Dict[str, Any]:
"""Get chunking statistics"""
stats = self.stats.copy()
# Calculate averages
if stats["contextual_chunks"] > 0:
stats["avg_context_tokens"] = stats["total_context_tokens"] / stats["contextual_chunks"]
stats["avg_generation_time"] = stats["total_generation_time"] / stats["contextual_chunks"]
else:
stats["avg_context_tokens"] = 0
stats["avg_generation_time"] = 0
# Cache efficiency
total_cache_ops = stats["cache_hits"] + stats["cache_misses"]
if total_cache_ops > 0:
stats["cache_hit_rate"] = stats["cache_hits"] / total_cache_ops
else:
stats["cache_hit_rate"] = 0
# Cost estimation
if self.llm_config.provider == "openai":
cost_per_1k = 0.03
else:
cost_per_1k = 0.01
stats["estimated_cost"] = (stats["total_context_tokens"] / 1000) * cost_per_1k
return stats
contextual_tools.py¶
"""Enhanced tools for contextual retrieval with BM25 and semantic search
Educational implementation showing how contextual chunks improve both
BM25 (lexical) and embedding (semantic) retrieval.
"""
import json
import logging
import requests
import numpy as np
from typing import Dict, Any, List, Optional, Tuple
from dataclasses import dataclass
from datetime import datetime
import time
from rank_bm25 import BM25Okapi
import pickle
from pathlib import Path
from config import KnowledgeBaseConfig, KnowledgeBaseType
from tools import KnowledgeBaseTools, SearchResult
from contextual_chunking import ContextualChunk
# Shared tokenizer: 中文没有空格，原先的 .lower().split() 会把整段当成一个 token，
# 导致 BM25 在中文语料上几乎失效。统一改用 compare_retrieval.tokenize（jieba 分词）。
from compare_retrieval import tokenize as _bm25_tokenize
logger = logging.getLogger(__name__)
@dataclass
class ContextualSearchResult(SearchResult):
"""Enhanced search result with contextual information"""
is_contextual: bool = False
context_text: str = ""
bm25_score: float = 0.0
embedding_score: float = 0.0
hybrid_score: float = 0.0
retrieval_method: str = "hybrid" # bm25, embedding, or hybrid
def to_dict(self) -> Dict[str, Any]:
base = super().to_dict()
base.update({
"is_contextual": self.is_contextual,
"context_text": self.context_text,
"bm25_score": self.bm25_score,
"embedding_score": self.embedding_score,
"hybrid_score": self.hybrid_score,
"retrieval_method": self.retrieval_method
})
return base
class ContextualKnowledgeBaseTools(KnowledgeBaseTools):
"""
Enhanced knowledge base tools with contextual retrieval support.
Key Educational Points:
1. Dual Indexing: Maintains both contextual and non-contextual indexes
2. BM25 Enhancement: Shows how context improves lexical matching
3. Hybrid Search: Combines BM25 and semantic search with rank fusion
4. Comparison Mode: Allows side-by-side evaluation of methods
"""
def __init__(self,
config: KnowledgeBaseConfig,
use_contextual: bool = True,
enable_comparison: bool = False):
"""
Initialize contextual knowledge base tools.
Args:
config: Knowledge base configuration
use_contextual: Whether to use contextual retrieval
enable_comparison: Whether to enable comparison mode
"""
super().__init__(config)
self.use_contextual = use_contextual
self.enable_comparison = enable_comparison
# BM25 indexes for lexical search
self.bm25_index = None
self.bm25_contextual_index = None
self.bm25_corpus = []
self.bm25_contextual_corpus = []
# Document and chunk storage
self.chunk_store = {} # chunk_id -> ContextualChunk
self.contextual_chunk_store = {} # chunk_id -> ContextualChunk (with context)
# Index paths
self.index_dir = Path("indexes")
self.index_dir.mkdir(exist_ok=True)
# Load existing indexes if available
self._load_indexes()
# Statistics
self.search_stats = {
"total_searches": 0,
"contextual_searches": 0,
"non_contextual_searches": 0,
"comparison_searches": 0,
"avg_retrieval_time": 0.0,
"total_retrieval_time": 0.0
}
logger.info(f"Initialized ContextualKnowledgeBaseTools (contextual={use_contextual}, comparison={enable_comparison})")
def index_contextual_chunks(self, chunks: List[ContextualChunk], rebuild_bm25: bool = True):
"""
Index contextual chunks for both BM25 and semantic search.
Educational Note:
This demonstrates the dual indexing strategy:
- BM25 index on contextualized text for better lexical matching
- Semantic embeddings on contextualized text for richer meaning
"""
logger.info(f"Indexing {len(chunks)} contextual chunks")
start_time = time.time()
# Store chunks
for chunk in chunks:
self.contextual_chunk_store[chunk.chunk_id] = chunk
# Also store non-contextual version for comparison
non_contextual_chunk = ContextualChunk(
chunk_id=chunk.chunk_id + "_nc",
doc_id=chunk.doc_id,
text=chunk.text,
context="",
contextualized_text=chunk.text,
chunk_index=chunk.chunk_index,
char_count=len(chunk.text),
metadata={"contextual": False}
)
self.chunk_store[non_contextual_chunk.chunk_id] = non_contextual_chunk
# Build BM25 indexes
if rebuild_bm25:
self._build_bm25_indexes()
# Index to retrieval pipeline (if local)
if self.config.type == KnowledgeBaseType.LOCAL:
self._index_chunks_to_pipeline(chunks)
elapsed = time.time() - start_time
logger.info(f"Indexed {len(chunks)} chunks in {elapsed:.2f}s")
# Save indexes
self._save_indexes()
def _build_bm25_indexes(self):
"""
Build BM25 indexes for both contextual and non-contextual chunks.
Educational Note:
BM25 uses TF-IDF with optimizations:
- Term frequency saturation prevents common words from dominating
- Document length normalization accounts for varying chunk sizes
- The contextual version has richer vocabulary from added context
"""
logger.info("Building BM25 indexes")
# Build contextual BM25 index
if self.contextual_chunk_store:
contextual_texts = []
for chunk in self.contextual_chunk_store.values():
# Tokenize for BM25 (jieba 中文分词，兼容英文)
tokens = _bm25_tokenize(chunk.contextualized_text)
contextual_texts.append(tokens)
self.bm25_contextual_corpus = contextual_texts
self.bm25_contextual_index = BM25Okapi(contextual_texts)
logger.info(f"Built contextual BM25 index with {len(contextual_texts)} documents")
# Build non-contextual BM25 index
if self.chunk_store:
non_contextual_texts = []
for chunk in self.chunk_store.values():
tokens = _bm25_tokenize(chunk.text)
non_contextual_texts.append(tokens)
self.bm25_corpus = non_contextual_texts
self.bm25_index = BM25Okapi(non_contextual_texts)
logger.info(f"Built non-contextual BM25 index with {len(non_contextual_texts)} documents")
def _index_chunks_to_pipeline(self, chunks: List[ContextualChunk]):
"""Index chunks to the retrieval pipeline"""
for chunk in chunks:
try:
# Index contextual version
if self.use_contextual:
response = requests.post(
f"{self.config.local_base_url}/index",
json={
"text": chunk.contextualized_text,
"doc_id": chunk.doc_id,
"metadata": {
"chunk_id": chunk.chunk_id,
"chunk_index": chunk.chunk_index,
"is_contextual": True,
"context": chunk.context[:200], # Store truncated context
"original_text": chunk.text[:500] # Store truncated original
}
}
)
response.raise_for_status()
# Also index non-contextual version if in comparison mode
if self.enable_comparison:
response = requests.post(
f"{self.config.local_base_url}/index",
json={
"text": chunk.text,
"doc_id": chunk.doc_id,
"metadata": {
"chunk_id": chunk.chunk_id + "_nc",
"chunk_index": chunk.chunk_index,
"is_contextual": False
}
}
)
response.raise_for_status()
except Exception as e:
logger.error(f"Error indexing chunk {chunk.chunk_id}: {e}")
def contextual_search(self,
query: str,
method: str = "hybrid",
top_k: int = 20) -> List[ContextualSearchResult]:
"""
Perform contextual search using specified method.
Args:
query: Search query
method: Search method - "bm25", "embedding", or "hybrid"
top_k: Number of results to return
Educational Note:
This demonstrates three retrieval strategies:
1. BM25: Pure lexical matching based on term frequency
2. Embedding: Semantic similarity using vector embeddings
3. Hybrid: Rank fusion combining both approaches
"""
logger.info(f"Performing {method} search for: {query[:100]}...")
start_time = time.time()
results = []
if method in ["bm25", "hybrid"]:
bm25_results = self._search_bm25(query, self.use_contextual, top_k * 2)
results.extend(bm25_results)
if method in ["embedding", "hybrid"]:
embedding_results = self._search_embeddings(query, self.use_contextual, top_k * 2)
results.extend(embedding_results)
if method == "hybrid":
# Rank fusion: combine and deduplicate results
results = self._rank_fusion(bm25_results, embedding_results, top_k)
else:
# Sort by score and limit
results = sorted(results, key=lambda x: x.score, reverse=True)[:top_k]
# Update statistics
elapsed = time.time() - start_time
self.search_stats["total_searches"] += 1
if self.use_contextual:
self.search_stats["contextual_searches"] += 1
else:
self.search_stats["non_contextual_searches"] += 1
self.search_stats["total_retrieval_time"] += elapsed
self.search_stats["avg_retrieval_time"] = (
self.search_stats["total_retrieval_time"] / self.search_stats["total_searches"]
)
logger.info(f"Search completed in {elapsed:.2f}s, returned {len(results)} results")
return results
def _search_bm25(self, query: str, use_contextual: bool, top_k: int) -> List[ContextualSearchResult]:
"""
Perform BM25 search.
Educational Note:
BM25 excels at finding exact term matches and handles
technical terms, IDs, and specific phrases well.
Contextual chunks help by adding synonyms and related terms.
"""
if use_contextual and self.bm25_contextual_index:
index = self.bm25_contextual_index
corpus = self.bm25_contextual_corpus
chunk_store = self.contextual_chunk_store
elif self.bm25_index:
index = self.bm25_index
corpus = self.bm25_corpus
chunk_store = self.chunk_store
else:
logger.warning("BM25 index not available")
return []
# Tokenize query (jieba 中文分词，兼容英文)
query_tokens = _bm25_tokenize(query)
# Get BM25 scores
scores = index.get_scores(query_tokens)
# Get top-k indices
top_indices = np.argsort(scores)[-top_k:][::-1]
# Create results
results = []
chunk_list = list(chunk_store.values())
for idx in top_indices:
if idx < len(chunk_list) and scores[idx] > 0:
chunk = chunk_list[idx]
result = ContextualSearchResult(
doc_id=chunk.doc_id,
chunk_id=chunk.chunk_id,
text=chunk.text,
score=float(scores[idx]),
is_contextual=use_contextual,
context_text=chunk.context if use_contextual else "",
bm25_score=float(scores[idx]),
retrieval_method="bm25",
metadata={"method": "bm25", "contextual": use_contextual}
)
results.append(result)
return results
def _search_embeddings(self, query: str, use_contextual: bool, top_k: int) -> List[ContextualSearchResult]:
"""
Perform semantic search using embeddings.
Educational Note:
Embedding search captures semantic meaning and relationships.
Contextual chunks provide richer semantic information,
helping find conceptually related content even without exact matches.
"""
try:
# Use the retrieval pipeline for embedding search
response = requests.post(
f"{self.config.local_base_url}/search",
json={
"query": query,
"mode": "embedding", # Use embedding mode
"top_k": top_k,
"filter": {"is_contextual": use_contextual} if self.enable_comparison else None
}
)
response.raise_for_status()
results = []
data = response.json()
for item in data.get("results", []):
# Map back to our chunk store
chunk_id = item.get("metadata", {}).get("chunk_id", "")
if use_contextual and chunk_id in self.contextual_chunk_store:
chunk = self.contextual_chunk_store[chunk_id]
elif chunk_id in self.chunk_store:
chunk = self.chunk_store[chunk_id]
else:
continue
result = ContextualSearchResult(
doc_id=chunk.doc_id,
chunk_id=chunk.chunk_id,
text=chunk.text,
score=item.get("score", 0.0),
is_contextual=use_contextual,
context_text=chunk.context if hasattr(chunk, 'context') else "",
embedding_score=item.get("score", 0.0),
retrieval_method="embedding",
metadata={"method": "embedding", "contextual": use_contextual}
)
results.append(result)
return results
except Exception as e:
logger.error(f"Error in embedding search: {e}")
return []
def _rank_fusion(self,
bm25_results: List[ContextualSearchResult],
embedding_results: List[ContextualSearchResult],
top_k: int) -> List[ContextualSearchResult]:
"""
Combine BM25 and embedding results using reciprocal rank fusion.
Educational Note:
Rank fusion combines different retrieval signals:
- BM25 provides strong exact matching
- Embeddings provide semantic understanding
- The combination often outperforms either method alone
We use Reciprocal Rank Fusion (RRF) which is simple but effective.
"""
fusion_scores = {}
chunk_map = {}
# RRF constant (typically 60)
k = 60
# Process BM25 results
for rank, result in enumerate(bm25_results):
rrf_score = 1.0 / (k + rank + 1)
fusion_scores[result.chunk_id] = fusion_scores.get(result.chunk_id, 0) + rrf_score
chunk_map[result.chunk_id] = result
result.bm25_score = result.score
# Process embedding results
for rank, result in enumerate(embedding_results):
rrf_score = 1.0 / (k + rank + 1)
if result.chunk_id in fusion_scores:
# Update existing result
fusion_scores[result.chunk_id] += rrf_score
chunk_map[result.chunk_id].embedding_score = result.score
else:
# New result from embeddings only
fusion_scores[result.chunk_id] = rrf_score
chunk_map[result.chunk_id] = result
result.embedding_score = result.score
# Create final results sorted by fusion score
final_results = []
for chunk_id, fusion_score in sorted(fusion_scores.items(),
key=lambda x: x[1],
reverse=True)[:top_k]:
result = chunk_map[chunk_id]
result.hybrid_score = fusion_score
result.score = fusion_score # Use fusion score as main score
result.retrieval_method = "hybrid"
final_results.append(result)
return final_results
def compare_retrieval_methods(self,
query: str,
top_k: int = 20) -> Dict[str, Any]:
"""
Compare contextual vs non-contextual retrieval.
Educational Note:
This method demonstrates the improvement that contextual
retrieval provides across different search methods.
It's useful for evaluation and understanding when context helps most.
"""
logger.info(f"Comparing retrieval methods for: {query[:100]}...")
comparison_results = {
"query": query,
"timestamp": datetime.now().isoformat(),
"top_k": top_k,
"methods": {}
}
# Test each combination
test_configs = [
("contextual_hybrid", True, "hybrid"),
("contextual_bm25", True, "bm25"),
("contextual_embedding", True, "embedding"),
("non_contextual_hybrid", False, "hybrid"),
("non_contextual_bm25", False, "bm25"),
("non_contextual_embedding", False, "embedding")
]
for name, use_contextual, method in test_configs:
# Temporarily set mode
original_contextual = self.use_contextual
self.use_contextual = use_contextual
# Perform search
start_time = time.time()
results = self.contextual_search(query, method, top_k)
elapsed = time.time() - start_time
# Store results
comparison_results["methods"][name] = {
"results": [r.to_dict() for r in results[:5]], # Top 5 for readability
"total_results": len(results),
"retrieval_time": elapsed,
"avg_score": np.mean([r.score for r in results]) if results else 0,
"max_score": max([r.score for r in results]) if results else 0
}
# Restore mode
self.use_contextual = original_contextual
# Add analysis
comparison_results["analysis"] = self._analyze_comparison(comparison_results)
# Update stats
self.search_stats["comparison_searches"] += 1
return comparison_results
def _analyze_comparison(self, results: Dict[str, Any]) -> Dict[str, Any]:
"""Analyze comparison results to highlight improvements"""
analysis = {
"contextual_improvement": {},
"method_comparison": {},
"recommendations": []
}
# Compare contextual vs non-contextual for each method
for method in ["hybrid", "bm25", "embedding"]:
contextual_key = f"contextual_{method}"
non_contextual_key = f"non_contextual_{method}"
if contextual_key in results["methods"] and non_contextual_key in results["methods"]:
contextual = results["methods"][contextual_key]
non_contextual = results["methods"][non_contextual_key]
# Calculate improvement
score_improvement = (
(contextual["avg_score"] - non_contextual["avg_score"])
/ non_contextual["avg_score"] * 100
if non_contextual["avg_score"] > 0 else 0
)
analysis["contextual_improvement"][method] = {
"score_improvement_pct": round(score_improvement, 2),
"contextual_avg_score": round(contextual["avg_score"], 4),
"non_contextual_avg_score": round(non_contextual["avg_score"], 4)
}
# Find best performing method
best_method = max(
results["methods"].items(),
key=lambda x: x[1]["avg_score"]
)
analysis["best_method"] = best_method[0]
# Generate recommendations
if "hybrid" in analysis["contextual_improvement"]:
if analysis["contextual_improvement"]["hybrid"]["score_improvement_pct"] > 10:
analysis["recommendations"].append(
"Contextual retrieval shows significant improvement (>10%). "
"Consider using it for production."
)
if analysis["contextual_improvement"]["bm25"]["score_improvement_pct"] > \
analysis["contextual_improvement"]["embedding"]["score_improvement_pct"]:
analysis["recommendations"].append(
"Contextual enhancement helps BM25 more than embeddings. "
"The query might contain specific terms that benefit from context."
)
return analysis
def _save_indexes(self):
"""Save BM25 indexes to disk"""
try:
# Save BM25 indexes
if self.bm25_contextual_index:
with open(self.index_dir / "bm25_contextual.pkl", "wb") as f:
pickle.dump({
"index": self.bm25_contextual_index,
"corpus": self.bm25_contextual_corpus
}, f)
if self.bm25_index:
with open(self.index_dir / "bm25_non_contextual.pkl", "wb") as f:
pickle.dump({
"index": self.bm25_index,
"corpus": self.bm25_corpus
}, f)
# Save chunk stores
with open(self.index_dir / "chunk_stores.json", "w") as f:
json.dump({
"contextual": {k: v.to_dict() for k, v in self.contextual_chunk_store.items()},
"non_contextual": {k: v.to_dict() for k, v in self.chunk_store.items()}
}, f, indent=2)
logger.info("Indexes saved successfully")
except Exception as e:
logger.error(f"Error saving indexes: {e}")
def _load_indexes(self):
"""Load BM25 indexes from disk"""
try:
# Load BM25 indexes
contextual_path = self.index_dir / "bm25_contextual.pkl"
if contextual_path.exists():
with open(contextual_path, "rb") as f:
data = pickle.load(f)
self.bm25_contextual_index = data["index"]
self.bm25_contextual_corpus = data["corpus"]
logger.info("Loaded contextual BM25 index")
non_contextual_path = self.index_dir / "bm25_non_contextual.pkl"
if non_contextual_path.exists():
with open(non_contextual_path, "rb") as f:
data = pickle.load(f)
self.bm25_index = data["index"]
self.bm25_corpus = data["corpus"]
logger.info("Loaded non-contextual BM25 index")
# Load chunk stores
stores_path = self.index_dir / "chunk_stores.json"
if stores_path.exists():
with open(stores_path, "r") as f:
data = json.load(f)
# Reconstruct contextual chunks
for chunk_id, chunk_dict in data.get("contextual", {}).items():
self.contextual_chunk_store[chunk_id] = ContextualChunk(
chunk_id=chunk_dict["chunk_id"],
doc_id=chunk_dict["doc_id"],
text=chunk_dict["text"],
context=chunk_dict["context"],
contextualized_text=chunk_dict["contextualized_text"],
chunk_index=chunk_dict["chunk_index"],
char_count=chunk_dict["char_count"],
context_tokens=chunk_dict.get("context_tokens", 0),
generation_time=chunk_dict.get("generation_time", 0),
metadata=chunk_dict.get("metadata", {})
)
# Reconstruct non-contextual chunks
for chunk_id, chunk_dict in data.get("non_contextual", {}).items():
self.chunk_store[chunk_id] = ContextualChunk(
chunk_id=chunk_dict["chunk_id"],
doc_id=chunk_dict["doc_id"],
text=chunk_dict["text"],
context="",
contextualized_text=chunk_dict["text"],
chunk_index=chunk_dict["chunk_index"],
char_count=chunk_dict["char_count"],
metadata=chunk_dict.get("metadata", {})
)
logger.info(f"Loaded {len(self.contextual_chunk_store)} contextual chunks")
logger.info(f"Loaded {len(self.chunk_store)} non-contextual chunks")
except Exception as e:
logger.info(f"No existing indexes found or error loading: {e}")
def get_statistics(self) -> Dict[str, Any]:
"""Get comprehensive statistics"""
stats = {
"search_stats": self.search_stats,
"index_stats": {
"contextual_chunks": len(self.contextual_chunk_store),
"non_contextual_chunks": len(self.chunk_store),
"bm25_contextual_size": len(self.bm25_contextual_corpus) if self.bm25_contextual_corpus else 0,
"bm25_non_contextual_size": len(self.bm25_corpus) if self.bm25_corpus else 0
}
}
return stats
demo.py¶
#!/usr/bin/env python3
"""Interactive demo of the Contextual Retrieval System
This script provides an interactive demonstration showing:
1. How chunks lose context in traditional RAG
2. How contextual retrieval solves this problem
3. Side-by-side comparison of retrieval quality
"""
import json
import logging
from pathlib import Path
from typing import List
import time
from datetime import datetime
from config import Config
from contextual_chunking import ContextualChunker
from contextual_tools import ContextualKnowledgeBaseTools
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(message)s' # Simple format for demo
)
logger = logging.getLogger(__name__)
def print_header(title: str, char: str = "=", width: int = 80):
"""Print a formatted header"""
logger.info(f"\n{char * width}")
logger.info(f"{title.center(width)}")
logger.info(f"{char * width}\n")
def print_section(title: str, char: str = "-", width: int = 60):
"""Print a section header"""
logger.info(f"\n{char * width}")
logger.info(f"{title}")
logger.info(f"{char * width}\n")
class ContextualRetrievalDemo:
"""Interactive demo class"""
def __init__(self):
self.config = Config.from_env()
self.documents = {}
self.contextual_kb = None
self.non_contextual_kb = None
def run(self):
"""Run the interactive demo"""
print_header("CONTEXTUAL RETRIEVAL SYSTEM - INTERACTIVE DEMO")
logger.info("Welcome! This demo will show you how contextual retrieval")
logger.info("improves RAG systems by preserving context when chunking.\n")
while True:
self.show_menu()
choice = input("\nYour choice: ").strip()
if choice == "1":
self.demo_problem()
elif choice == "2":
self.demo_solution()
elif choice == "3":
self.demo_comparison()
elif choice == "4":
self.demo_real_example()
elif choice == "5":
self.show_statistics()
elif choice == "q":
logger.info("\nThank you for using the Contextual Retrieval Demo!")
break
else:
logger.info("Invalid choice. Please try again.")
def show_menu(self):
"""Show main menu"""
print_section("MAIN MENU")
logger.info("1. Demonstrate the Context Loss Problem")
logger.info("2. Show the Contextual Retrieval Solution")
logger.info("3. Compare Search Results (Side-by-Side)")
logger.info("4. Real-World Example (Financial Report)")
logger.info("5. Show Performance Statistics")
logger.info("q. Quit")
def demo_problem(self):
"""Demonstrate the context loss problem"""
print_header("THE CONTEXT LOSS PROBLEM", char="*")
# Example document
document = """
ACME Corporation Annual Report 2023
Financial Highlights:
ACME Corporation achieved record performance in 2023. The company's revenue
grew by 15% compared to the previous year, reaching $2.5 billion. This growth
was driven by strong demand in the technology sector.
TechStart Inc. Performance:
Meanwhile, TechStart Inc. faced challenges in 2023. The company's revenue
declined by 8% due to supply chain disruptions. Management has implemented
cost-cutting measures to improve profitability.
Global Industries Update:
Global Industries maintained steady performance. The company's revenue
remained flat at $1.8 billion, but profit margins improved by 2 percentage
points through operational efficiency gains.
""".strip()
logger.info("Consider this document:\n")
logger.info("=" * 60)
logger.info(document)
logger.info("=" * 60)
logger.info("\nNow imagine we chunk this document and get:\n")
# Show problematic chunks
chunks = [
"The company's revenue grew by 15% compared to the previous year, reaching $2.5 billion.",
"The company's revenue declined by 8% due to supply chain disruptions.",
"The company's revenue remained flat at $1.8 billion, but profit margins improved."
]
for i, chunk in enumerate(chunks, 1):
logger.info(f"CHUNK {i}:")
logger.info(f" '{chunk}'")
logger.info("")
logger.info("❌ PROBLEM: All chunks say 'The company' but refer to different companies!")
logger.info("❌ A search for 'company revenue growth' might return the wrong chunk!")
logger.info("❌ Without context, we can't tell which company each chunk refers to!\n")
input("\nPress Enter to continue...")
def demo_solution(self):
"""Demonstrate the contextual retrieval solution"""
print_header("THE CONTEXTUAL RETRIEVAL SOLUTION", char="*")
logger.info("Contextual Retrieval solves this by adding context to each chunk:\n")
# Show contextualized chunks
contextual_chunks = [
{
"context": "This chunk is from ACME Corporation's 2023 annual report financial highlights section.",
"text": "The company's revenue grew by 15% compared to the previous year, reaching $2.5 billion."
},
{
"context": "This chunk discusses TechStart Inc.'s 2023 performance challenges.",
"text": "The company's revenue declined by 8% due to supply chain disruptions."
},
{
"context": "This chunk covers Global Industries' steady 2023 performance.",
"text": "The company's revenue remained flat at $1.8 billion, but profit margins improved."
}
]
for i, chunk in enumerate(contextual_chunks, 1):
logger.info(f"CONTEXTUAL CHUNK {i}:")
logger.info(f" Context: {chunk['context']}")
logger.info(f" Text: {chunk['text']}")
logger.info(f" Combined: {chunk['context']} {chunk['text']}\n")
logger.info("✅ SOLUTION: Each chunk now has context!")
logger.info("✅ Searching for 'ACME revenue growth' will correctly find chunk 1!")
logger.info("✅ The context preserves crucial information lost in traditional chunking!\n")
input("\nPress Enter to continue...")
def demo_comparison(self):
"""Run a side-by-side comparison"""
print_header("SIDE-BY-SIDE COMPARISON", char="*")
# Create test document
test_doc = """
Artificial Intelligence in Healthcare
Introduction:
Artificial intelligence is transforming healthcare delivery. Machine learning
models are being used for disease diagnosis, drug discovery, and patient care
optimization. The technology has shown remarkable results in early detection
of diseases.
Diagnostic Applications:
In radiology, AI systems can detect cancer with 95% accuracy. The systems
analyze medical images faster than human radiologists. This reduces diagnosis
time from hours to minutes.
Drug Discovery:
Pharmaceutical companies use AI to identify potential drug compounds. The
technology can predict drug interactions and side effects. This accelerates
the drug development process by years.
""".strip()
logger.info("Test Document:")
logger.info("=" * 60)
logger.info(test_doc[:300] + "..." if len(test_doc) > 300 else test_doc)
logger.info("=" * 60)
# Initialize systems
logger.info("\nInitializing systems...")
# Create contextual system
contextual_chunker = ContextualChunker(
chunking_config=self.config.chunking,
llm_config=self.config.llm,
use_contextual=True
)
self.contextual_kb = ContextualKnowledgeBaseTools(
config=self.config.knowledge_base,
use_contextual=True
)
# Create non-contextual system
non_contextual_chunker = ContextualChunker(
chunking_config=self.config.chunking,
llm_config=self.config.llm,
use_contextual=False
)
self.non_contextual_kb = ContextualKnowledgeBaseTools(
config=self.config.knowledge_base,
use_contextual=False
)
# Process document
logger.info("\nProcessing document...")
# Contextual chunks
contextual_chunks = contextual_chunker.chunk_document(
text=test_doc,
doc_id="healthcare_ai"
)
self.contextual_kb.index_contextual_chunks(contextual_chunks)
# Non-contextual chunks
non_contextual_chunks = non_contextual_chunker.chunk_document(
text=test_doc,
doc_id="healthcare_ai"
)
self.non_contextual_kb.index_contextual_chunks(non_contextual_chunks)
# Test queries
test_queries = [
"How accurate is the AI system?",
"What technology reduces diagnosis time?",
"What can the technology predict?"
]
logger.info("\nRunning comparison...")
for query in test_queries:
print_section(f"Query: {query}")
# Contextual search
contextual_results = self.contextual_kb.contextual_search(
query=query,
method="hybrid",
top_k=3
)
# Non-contextual search
non_contextual_results = self.non_contextual_kb.contextual_search(
query=query,
method="hybrid",
top_k=3
)
logger.info("CONTEXTUAL RESULTS:")
if contextual_results:
result = contextual_results[0]
logger.info(f" Score: {result.score:.4f}")
logger.info(f" Context: {result.context_text[:100]}..." if result.context_text else " Context: None")
logger.info(f" Match: {result.text[:100]}...\n")
else:
logger.info(" No results found\n")
logger.info("NON-CONTEXTUAL RESULTS:")
if non_contextual_results:
result = non_contextual_results[0]
logger.info(f" Score: {result.score:.4f}")
logger.info(f" Match: {result.text[:100]}...\n")
else:
logger.info(" No results found\n")
# Score comparison
if contextual_results and non_contextual_results:
improvement = ((contextual_results[0].score - non_contextual_results[0].score)
/ non_contextual_results[0].score * 100)
logger.info(f"📊 Score Improvement: {improvement:.1f}%\n")
input("\nPress Enter to continue...")
def demo_real_example(self):
"""Demonstrate with a real-world example"""
print_header("REAL-WORLD EXAMPLE: FINANCIAL REPORT", char="*")
# Create a realistic financial report
report = """
Q2 2023 Earnings Report - TechCorp International
Executive Summary:
TechCorp International reported strong second quarter results for 2023,
with revenue of $850 million, representing a 12% year-over-year growth.
The company's cloud services division was the primary growth driver.
Revenue Breakdown:
Cloud Services: Revenue increased by 25% to $400 million, driven by
enterprise adoption of our AI-powered analytics platform. Operating margin
improved to 35% from 30% in the prior year.
Hardware Division: Revenue declined by 5% to $300 million due to supply
chain constraints. However, the new product pipeline remains strong with
three launches planned for Q3.
Software Licensing: Revenue grew by 8% to $150 million. The company added
200 new enterprise customers during the quarter, bringing the total to
5,000 active licenses.
Competitive Analysis:
Compared to DataSoft Corp, our main competitor, we maintained market share
leadership. DataSoft reported 8% revenue growth in their latest quarter,
while our 12% growth demonstrates strong execution. Their cloud division
grew 15% compared to our 25% growth.
Future Outlook:
Management expects continued momentum in Q3 2023. The company raised
full-year guidance to $3.5 billion in revenue, representing 15% annual
growth. Investment in R&D will increase by 20% to accelerate AI
product development.
""".strip()
logger.info("Processing a realistic financial report...")
logger.info("=" * 60)
logger.info(report[:400] + "..." if len(report) > 400 else report)
logger.info("=" * 60)
# Process with contextual system
logger.info("\nGenerating contextual chunks (this may take a moment)...\n")
chunker = ContextualChunker(
chunking_config=self.config.chunking,
llm_config=self.config.llm,
use_contextual=True
)
chunks = chunker.chunk_document(
text=report,
doc_id="techcorp_q2_2023"
)
# Show examples of contextualized chunks
logger.info("Example Contextual Chunks:\n")
for chunk in chunks[:3]:
logger.info(f"Original: {chunk.text[:100]}...")
logger.info(f"Context: {chunk.context}")
logger.info("")
# Show how this helps with ambiguous queries
logger.info("Why this matters:\n")
logger.info("Query: 'What was the revenue growth?'")
logger.info(" - Without context: Could match TechCorp's 12%, Cloud's 25%, or Software's 8%")
logger.info(" - With context: Correctly identifies which growth figure you want\n")
logger.info("Query: 'How did the company perform vs competition?'")
logger.info(" - Without context: Might return DataSoft's results")
logger.info(" - With context: Returns TechCorp's performance comparison\n")
input("\nPress Enter to continue...")
def show_statistics(self):
"""Show performance statistics"""
print_header("PERFORMANCE STATISTICS", char="*")
if not self.contextual_kb:
logger.info("No searches performed yet. Run a comparison first!\n")
input("\nPress Enter to continue...")
return
# Get statistics
stats = self.contextual_kb.get_statistics()
logger.info("Search Statistics:")
logger.info(f" Total searches: {stats['search_stats']['total_searches']}")
logger.info(f" Average retrieval time: {stats['search_stats']['avg_retrieval_time']:.3f}s")
logger.info(f" Contextual searches: {stats['search_stats']['contextual_searches']}")
logger.info(f" Non-contextual searches: {stats['search_stats']['non_contextual_searches']}")
logger.info("\nIndex Statistics:")
logger.info(f" Contextual chunks indexed: {stats['index_stats']['contextual_chunks']}")
logger.info(f" Non-contextual chunks indexed: {stats['index_stats']['non_contextual_chunks']}")
logger.info("\nBased on Anthropic's Research:")
logger.info(" Standard RAG: 5.7% retrieval failure rate")
logger.info(" Contextual RAG: 2.9% failure rate (49% improvement)")
logger.info(" + Reranking: 1.9% failure rate (67% improvement)")
input("\nPress Enter to continue...")
def main():
"""Run the interactive demo"""
demo = ContextualRetrievalDemo()
demo.run()
if __name__ == "__main__":
main()
evaluation/dataset_builder.py¶
"""Build evaluation dataset from Chinese legal documents"""
import json
import logging
from typing import List, Dict, Any
from pathlib import Path
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class LegalDatasetBuilder:
"""Build evaluation dataset for Chinese legal Q&A"""
def __init__(self):
self.simple_cases = []
self.complex_cases = []
def create_simple_cases(self) -> List[Dict[str, Any]]:
"""Create simple direct legal questions"""
simple_cases = [
{
"id": "simple_1",
"question": "故意杀人罪判几年？",
"expected_keywords": ["死刑", "无期徒刑", "十年以上有期徒刑"],
"reference": "《中华人民共和国刑法》第二百三十二条",
"difficulty": "easy"
},
{
"id": "simple_2",
"question": "盗窃罪的立案标准是什么？",
"expected_keywords": ["一千元", "三千元", "数额较大"],
"reference": "《中华人民共和国刑法》第二百六十四条",
"difficulty": "easy"
},
{
"id": "simple_3",
"question": "醉酒驾驶机动车如何处罚？",
"expected_keywords": ["拘役", "罚金", "吊销驾照"],
"reference": "《中华人民共和国刑法》第一百三十三条",
"difficulty": "easy"
},
{
"id": "simple_4",
"question": "诈骗罪的量刑标准是什么？",
"expected_keywords": ["三年以下", "三年以上十年以下", "十年以上"],
"reference": "《中华人民共和国刑法》第二百六十六条",
"difficulty": "easy"
},
{
"id": "simple_5",
"question": "故意伤害罪致人重伤的处罚是什么？",
"expected_keywords": ["三年以上十年以下", "有期徒刑"],
"reference": "《中华人民共和国刑法》第二百三十四条",
"difficulty": "easy"
},
{
"id": "simple_6",
"question": "抢劫罪的加重情节有哪些？",
"expected_keywords": ["入户抢劫", "多次抢劫", "抢劫数额巨大"],
"reference": "《中华人民共和国刑法》第二百六十三条",
"difficulty": "medium"
},
{
"id": "simple_7",
"question": "非法拘禁罪的构成要件是什么？",
"expected_keywords": ["非法", "拘禁", "限制人身自由"],
"reference": "《中华人民共和国刑法》第二百三十八条",
"difficulty": "medium"
},
{
"id": "simple_8",
"question": "贪污罪的数额标准如何认定？",
"expected_keywords": ["三万元", "二十万元", "三百万元"],
"reference": "《中华人民共和国刑法》第三百八十三条",
"difficulty": "medium"
},
{
"id": "simple_9",
"question": "交通肇事罪的立案标准是什么？",
"expected_keywords": ["死亡一人", "重伤三人", "财产损失"],
"reference": "《中华人民共和国刑法》第一百三十三条",
"difficulty": "easy"
},
{
"id": "simple_10",
"question": "寻衅滋事罪如何处罚？",
"expected_keywords": ["五年以下", "有期徒刑", "拘役", "管制"],
"reference": "《中华人民共和国刑法》第二百九十三条",
"difficulty": "easy"
}
]
return simple_cases
def create_complex_cases(self) -> List[Dict[str, Any]]:
"""Create complex legal scenario questions"""
complex_cases = [
{
"id": "complex_1",
"question": """张某因与李某发生经济纠纷，持刀闯入李某家中，意图讨债。在争执过程中，张某用刀刺伤李某，
导致李某重伤。同时，张某还顺手拿走了李某家中的现金5万元。请问张某的行为应如何定性？
可能面临什么样的刑事处罚？""",
"expected_analysis": ["入户抢劫", "故意伤害", "数罪并罚"],
"reference": "《刑法》第二百三十四条、第二百六十三条",
"difficulty": "hard",
"requires_multi_query": True
},
{
"id": "complex_2",
"question": """王某系某国有企业财务主管，利用职务之便，通过虚开发票等手段，
将公司资金200万元转入其控制的账户。后王某用该资金进行股票投资，
获利50万元。案发后，王某主动退还全部赃款。请分析王某的法律责任。""",
"expected_analysis": ["贪污罪", "挪用公款罪", "自首情节", "退赃"],
"reference": "《刑法》第三百八十二条、第三百八十四条",
"difficulty": "hard",
"requires_multi_query": True
},
{
"id": "complex_3",
"question": """赵某酒后驾车，在市区超速行驶，撞倒正在过马路的行人陈某，
导致陈某当场死亡。赵某见状，驾车逃离现场。第二天，在家人劝说下，
赵某到公安机关投案自首。请问赵某涉嫌哪些犯罪？量刑时应考虑哪些因素？""",
"expected_analysis": ["交通肇事罪", "危险驾驶罪", "逃逸", "自首"],
"reference": "《刑法》第一百三十三条",
"difficulty": "hard",
"requires_multi_query": True
},
{
"id": "complex_4",
"question": """刘某通过网络平台发布虚假投资信息，声称可以保证高额回报，
先后骗取30名投资者共计500万元。其中，刘某将200万元用于个人挥霍，
300万元用于归还之前的债务。请问刘某的行为如何定性？可能的量刑是什么？""",
"expected_analysis": ["诈骗罪", "数额特别巨大", "多人受害"],
"reference": "《刑法》第二百六十六条",
"difficulty": "hard",
"requires_multi_query": True
},
{
"id": "complex_5",
"question": """孙某与钱某共谋盗窃某商场。孙某负责望风，钱某进入商场实施盗窃。
钱某在盗窃过程中被保安发现，为逃跑将保安打成轻伤。
最终二人盗窃财物价值8万元。请分析孙某和钱某各自的刑事责任。""",
"expected_analysis": ["共同犯罪", "盗窃罪", "抢劫罪", "转化犯"],
"reference": "《刑法》第二百六十四条、第二百六十九条",
"difficulty": "hard",
"requires_multi_query": True
}
]
return complex_cases
def build_dataset(self, output_path: str = "legal_qa_dataset.json"):
"""Build and save the complete dataset"""
dataset = {
"simple_cases": self.create_simple_cases(),
"complex_cases": self.create_complex_cases(),
"metadata": {
"total_cases": 15,
"simple_count": 10,
"complex_count": 5,
"domain": "Chinese Criminal Law",
"purpose": "Evaluate agentic vs non-agentic RAG performance"
}
}
# Save dataset
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(dataset, f, ensure_ascii=False, indent=2)
logger.info(f"Dataset saved to {output_path}")
return dataset
def create_legal_documents() -> List[Dict[str, str]]:
"""Create sample legal documents for the knowledge base"""
documents = [
{
"doc_id": "criminal_law_homicide",
"title": "刑法-故意杀人罪",
"content": """第二百三十二条 【故意杀人罪】故意杀人的，处死刑、无期徒刑或者十年以上有期徒刑；
情节较轻的，处三年以上十年以下有期徒刑。
故意杀人罪是指故意非法剥夺他人生命的行为。该罪侵犯的客体是他人的生命权。
法律依据是《中华人民共和国刑法》第二百三十二条。
量刑标准：
1. 情节严重的：死刑、无期徒刑或十年以上有期徒刑
2. 情节较轻的：三年以上十年以下有期徒刑
情节较轻通常包括：防卫过当、义愤杀人、被害人有过错等情形。"""
},
{
"doc_id": "criminal_law_theft",
"title": "刑法-盗窃罪",
"content": """第二百六十四条 【盗窃罪】盗窃公私财物，数额较大的，或者多次盗窃、入户盗窃、
携带凶器盗窃、扒窃的，处三年以下有期徒刑、拘役或者管制，并处或者单处罚金；
数额巨大或者有其他严重情节的，处三年以上十年以下有期徒刑，并处罚金；
数额特别巨大或者有其他特别严重情节的，处十年以上有期徒刑或者无期徒刑，并处罚金或者没收财产。
盗窃罪的立案标准：
1. 数额较大：一般为1000元至3000元以上
2. 数额巨大：一般为3万元至10万元以上
3. 数额特别巨大：一般为30万元至50万元以上
特殊情形：多次盗窃（2年内3次以上）、入户盗窃、携带凶器盗窃、扒窃的，
不论数额大小，均构成盗窃罪。"""
},
{
"doc_id": "criminal_law_fraud",
"title": "刑法-诈骗罪",
"content": """第二百六十六条 【诈骗罪】诈骗公私财物，数额较大的，处三年以下有期徒刑、
拘役或者管制，并处或者单处罚金；数额巨大或者有其他严重情节的，
处三年以上十年以下有期徒刑，并处罚金；数额特别巨大或者有其他特别严重情节的，
处十年以上有期徒刑或者无期徒刑，并处罚金或者没收财产。
诈骗罪的量刑标准：
1. 数额较大（3千元至1万元以上）：三年以下有期徒刑、拘役或者管制
2. 数额巨大（3万元至10万元以上）：三年以上十年以下有期徒刑
3. 数额特别巨大（50万元以上）：十年以上有期徒刑或者无期徒刑
诈骗罪是指以非法占有为目的，用虚构事实或者隐瞒真相的方法，
骗取数额较大的公私财物的行为。"""
},
{
"doc_id": "criminal_law_robbery",
"title": "刑法-抢劫罪",
"content": """第二百六十三条 【抢劫罪】以暴力、胁迫或者其他方法抢劫公私财物的，
处三年以上十年以下有期徒刑，并处罚金；有下列情形之一的，
处十年以上有期徒刑、无期徒刑或者死刑，并处罚金或者没收财产：
（一）入户抢劫的；
（二）在公共交通工具上抢劫的；
（三）抢劫银行或者其他金融机构的；
（四）多次抢劫或者抢劫数额巨大的；
（五）抢劫致人重伤、死亡的；
（六）冒充军警人员抢劫的；
（七）持枪抢劫的；
（八）抢劫军用物资或者抢险、救灾、救济物资的。
抢劫罪的加重处罚情节包括上述八种情形，有其中之一的，
最低刑期为十年有期徒刑。"""
},
{
"doc_id": "criminal_law_injury",
"title": "刑法-故意伤害罪",
"content": """第二百三十四条 【故意伤害罪】故意伤害他人身体的，处三年以下有期徒刑、
拘役或者管制。犯前款罪，致人重伤的，处三年以上十年以下有期徒刑；
致人死亡或者以特别残忍手段致人重伤造成严重残疾的，处十年以上有期徒刑、
无期徒刑或者死刑。
故意伤害罪的量刑：
1. 故意伤害致人轻伤的：三年以下有期徒刑、拘役或者管制
2. 故意伤害致人重伤的：三年以上十年以下有期徒刑
3. 故意伤害致人死亡或特别残忍手段致残的：十年以上有期徒刑、无期徒刑或死刑
重伤标准：使人肢体残废或者毁人容貌；使人丧失听觉、视觉或者其他器官功能；
其他对于人身健康有重大伤害的。"""
},
{
"doc_id": "criminal_law_traffic",
"title": "刑法-交通肇事罪与危险驾驶罪",
"content": """第一百三十三条 【交通肇事罪】违反交通运输管理法规，因而发生重大事故，
致人重伤、死亡或者使公私财产遭受重大损失的，处三年以下有期徒刑或者拘役；
交通运输肇事后逃逸或者有其他特别恶劣情节的，处三年以上七年以下有期徒刑；
因逃逸致人死亡的，处七年以上有期徒刑。
第一百三十三条之一 【危险驾驶罪】在道路上驾驶机动车，有下列情形之一的，
处拘役，并处罚金：
（一）追逐竞驶，情节恶劣的；
（二）醉酒驾驶机动车的；
（三）从事校车业务或者旅客运输，严重超过额定乘员载客，
或者严重超过规定时速行驶的；
（四）违反危险化学品安全管理规定运输危险化学品，危及公共安全的。
醉酒驾驶的认定标准：血液酒精含量达到80毫克/100毫升以上。"""
},
{
"doc_id": "criminal_law_corruption",
"title": "刑法-贪污罪",
"content": """第三百八十二条 【贪污罪】国家工作人员利用职务上的便利，侵吞、窃取、
骗取或者以其他手段非法占有公共财物的，是贪污罪。
第三百八十三条 【贪污罪的处罚规定】对犯贪污罪的，根据情节轻重，分别依照下列规定处罚：
（一）贪污数额较大或者有其他较重情节的，处三年以下有期徒刑或者拘役，并处罚金。
（二）贪污数额巨大或者有其他严重情节的，处三年以上十年以下有期徒刑，并处罚金或者没收财产。
（三）贪污数额特别巨大或者有其他特别严重情节的，处十年以上有期徒刑或者无期徒刑，
并处罚金或者没收财产；数额特别巨大，并使国家和人民利益遭受特别重大损失的，
处无期徒刑或者死刑，并处没收财产。
贪污数额标准：
1. 数额较大：三万元以上不满二十万元
2. 数额巨大：二十万元以上不满三百万元
3. 数额特别巨大：三百万元以上"""
}
]
return documents
if __name__ == "__main__":
# Build evaluation dataset
builder = LegalDatasetBuilder()
dataset = builder.build_dataset("legal_qa_dataset.json")
print(f"Dataset created with {len(dataset['simple_cases'])} simple cases and {len(dataset['complex_cases'])} complex cases")
# Create legal documents
documents = create_legal_documents()
# Save documents
with open("legal_documents.json", 'w', encoding='utf-8') as f:
json.dump(documents, f, ensure_ascii=False, indent=2)
print(f"Created {len(documents)} legal documents for knowledge base")
evaluation/evaluate.py¶
"""Evaluation framework for Agentic RAG system"""
import json
import logging
import time
from typing import List, Dict, Any, Optional
from pathlib import Path
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import Config
from agent import AgenticRAG
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class RAGEvaluator:
"""Evaluate RAG system performance"""
def __init__(self, config: Optional[Config] = None):
self.config = config or Config.from_env()
self.agent = AgenticRAG(self.config)
self.results = {
"agentic": [],
"non_agentic": []
}
def load_dataset(self, dataset_path: str) -> Dict[str, Any]:
"""Load evaluation dataset"""
with open(dataset_path, 'r', encoding='utf-8') as f:
return json.load(f)
def evaluate_response(self,
response: str,
test_case: Dict[str, Any]) -> Dict[str, Any]:
"""Evaluate a single response"""
evaluation = {
"case_id": test_case["id"],
"question": test_case["question"],
"response": response,
"metrics": {}
}
# Check for expected keywords (for simple cases)
if "expected_keywords" in test_case:
keywords_found = []
keywords_missing = []
for keyword in test_case["expected_keywords"]:
if keyword.lower() in response.lower():
keywords_found.append(keyword)
else:
keywords_missing.append(keyword)
evaluation["metrics"]["keyword_recall"] = len(keywords_found) / len(test_case["expected_keywords"])
evaluation["metrics"]["keywords_found"] = keywords_found
evaluation["metrics"]["keywords_missing"] = keywords_missing
# Check for analysis points (for complex cases)
if "expected_analysis" in test_case:
analysis_found = []
analysis_missing = []
for point in test_case["expected_analysis"]:
if point.lower() in response.lower():
analysis_found.append(point)
else:
analysis_missing.append(point)
evaluation["metrics"]["analysis_recall"] = len(analysis_found) / len(test_case["expected_analysis"])
evaluation["metrics"]["analysis_found"] = analysis_found
evaluation["metrics"]["analysis_missing"] = analysis_missing
# Check for citations
citation_count = response.count("[Doc:") + response.count("[Chunk:")
evaluation["metrics"]["has_citations"] = citation_count > 0
evaluation["metrics"]["citation_count"] = citation_count
# Response length
evaluation["metrics"]["response_length"] = len(response)
# Check if response indicates no answer
no_answer_indicators = ["无法回答", "没有找到", "知识库中没有", "cannot answer", "not found"]
evaluation["metrics"]["gave_answer"] = not any(indicator in response.lower() for indicator in no_answer_indicators)
return evaluation
def run_test_case(self, test_case: Dict[str, Any], mode: str = "agentic") -> Dict[str, Any]:
"""Run a single test case"""
logger.info(f"Running {mode} mode for case {test_case['id']}")
start_time = time.time()
try:
if mode == "agentic":
response = self.agent.query(test_case["question"], stream=False)
else:
response = self.agent.query_non_agentic(test_case["question"], stream=False)
elapsed_time = time.time() - start_time
# Clear history for next test
self.agent.clear_history()
# Evaluate response
evaluation = self.evaluate_response(response, test_case)
evaluation["mode"] = mode
evaluation["elapsed_time"] = elapsed_time
evaluation["difficulty"] = test_case.get("difficulty", "unknown")
evaluation["success"] = True
except Exception as e:
logger.error(f"Error in test case {test_case['id']}: {e}")
evaluation = {
"case_id": test_case["id"],
"question": test_case["question"],
"mode": mode,
"success": False,
"error": str(e),
"elapsed_time": time.time() - start_time
}
return evaluation
def run_evaluation(self, dataset_path: str, output_dir: str = "results"):
"""Run full evaluation"""
# Load dataset
dataset = self.load_dataset(dataset_path)
# Create output directory
output_path = Path(output_dir)
output_path.mkdir(exist_ok=True)
# Combine all test cases
all_cases = dataset["simple_cases"] + dataset["complex_cases"]
# Run agentic mode
logger.info("=" * 60)
logger.info("Running AGENTIC mode evaluation")
logger.info("=" * 60)
agentic_results = []
for test_case in all_cases:
result = self.run_test_case(test_case, mode="agentic")
agentic_results.append(result)
time.sleep(1) # Rate limiting
# Run non-agentic mode
logger.info("=" * 60)
logger.info("Running NON-AGENTIC mode evaluation")
logger.info("=" * 60)
non_agentic_results = []
for test_case in all_cases:
result = self.run_test_case(test_case, mode="non_agentic")
non_agentic_results.append(result)
time.sleep(1) # Rate limiting
# Compute aggregate metrics
agentic_metrics = self.compute_aggregate_metrics(agentic_results)
non_agentic_metrics = self.compute_aggregate_metrics(non_agentic_results)
# Save results
results = {
"dataset": dataset_path,
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
"config": {
"llm_provider": self.config.llm.provider,
"llm_model": self.agent.model,
"kb_type": self.config.knowledge_base.type.value
},
"agentic": {
"results": agentic_results,
"metrics": agentic_metrics
},
"non_agentic": {
"results": non_agentic_results,
"metrics": non_agentic_metrics
},
"comparison": self.compare_modes(agentic_metrics, non_agentic_metrics)
}
# Save to file
output_file = output_path / f"evaluation_results_{time.strftime('%Y%m%d_%H%M%S')}.json"
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(results, f, ensure_ascii=False, indent=2)
logger.info(f"Results saved to {output_file}")
# Print summary
self.print_summary(results)
return results
def compute_aggregate_metrics(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
"""Compute aggregate metrics from results"""
metrics = {
"total_cases": len(results),
"successful_cases": sum(1 for r in results if r.get("success", False)),
"failed_cases": sum(1 for r in results if not r.get("success", False)),
"average_time": 0,
"total_time": 0
}
# Separate by difficulty
simple_results = [r for r in results if r.get("difficulty") == "easy"]
medium_results = [r for r in results if r.get("difficulty") == "medium"]
hard_results = [r for r in results if r.get("difficulty") == "hard"]
# Compute metrics for successful cases
successful_results = [r for r in results if r.get("success", False)]
if successful_results:
# Time metrics
times = [r["elapsed_time"] for r in successful_results]
metrics["average_time"] = sum(times) / len(times)
metrics["total_time"] = sum(times)
metrics["min_time"] = min(times)
metrics["max_time"] = max(times)
# Response quality metrics
metrics["cases_with_citations"] = sum(1 for r in successful_results
if r.get("metrics", {}).get("has_citations", False))
metrics["cases_gave_answer"] = sum(1 for r in successful_results
if r.get("metrics", {}).get("gave_answer", False))
# Average response length
lengths = [r.get("metrics", {}).get("response_length", 0) for r in successful_results]
metrics["average_response_length"] = sum(lengths) / len(lengths) if lengths else 0
# Keyword/analysis recall (for cases that have them)
keyword_recalls = [r["metrics"]["keyword_recall"] for r in successful_results
if "keyword_recall" in r.get("metrics", {})]
if keyword_recalls:
metrics["average_keyword_recall"] = sum(keyword_recalls) / len(keyword_recalls)
analysis_recalls = [r["metrics"]["analysis_recall"] for r in successful_results
if "analysis_recall" in r.get("metrics", {})]
if analysis_recalls:
metrics["average_analysis_recall"] = sum(analysis_recalls) / len(analysis_recalls)
# Metrics by difficulty
for difficulty, diff_results in [("easy", simple_results), ("medium", medium_results), ("hard", hard_results)]:
if diff_results:
successful = [r for r in diff_results if r.get("success", False)]
metrics[f"{difficulty}_success_rate"] = len(successful) / len(diff_results)
if successful:
times = [r["elapsed_time"] for r in successful]
metrics[f"{difficulty}_average_time"] = sum(times) / len(times)
return metrics
def compare_modes(self, agentic_metrics: Dict[str, Any], non_agentic_metrics: Dict[str, Any]) -> Dict[str, Any]:
"""Compare agentic vs non-agentic performance"""
comparison = {}
# Success rate comparison
comparison["success_rate_diff"] = (agentic_metrics.get("successful_cases", 0) / agentic_metrics["total_cases"] -
non_agentic_metrics.get("successful_cases", 0) / non_agentic_metrics["total_cases"])
# Time comparison
if "average_time" in agentic_metrics and "average_time" in non_agentic_metrics:
comparison["time_ratio"] = agentic_metrics["average_time"] / non_agentic_metrics["average_time"]
comparison["time_difference"] = agentic_metrics["average_time"] - non_agentic_metrics["average_time"]
# Citation comparison
if "cases_with_citations" in agentic_metrics and "cases_with_citations" in non_agentic_metrics:
comparison["citation_rate_diff"] = (agentic_metrics["cases_with_citations"] / agentic_metrics["successful_cases"] -
non_agentic_metrics["cases_with_citations"] / non_agentic_metrics["successful_cases"])
# Response quality comparison
if "average_keyword_recall" in agentic_metrics and "average_keyword_recall" in non_agentic_metrics:
comparison["keyword_recall_improvement"] = (agentic_metrics["average_keyword_recall"] -
non_agentic_metrics["average_keyword_recall"])
if "average_analysis_recall" in agentic_metrics and "average_analysis_recall" in non_agentic_metrics:
comparison["analysis_recall_improvement"] = (agentic_metrics["average_analysis_recall"] -
non_agentic_metrics["average_analysis_recall"])
# Difficulty-specific comparison
for difficulty in ["easy", "medium", "hard"]:
key = f"{difficulty}_success_rate"
if key in agentic_metrics and key in non_agentic_metrics:
comparison[f"{difficulty}_success_improvement"] = (agentic_metrics[key] - non_agentic_metrics[key])
return comparison
def print_summary(self, results: Dict[str, Any]):
"""Print evaluation summary"""
print("\n" + "=" * 80)
print("EVALUATION SUMMARY")
print("=" * 80)
print(f"\nConfiguration:")
print(f" LLM Provider: {results['config']['llm_provider']}")
print(f" LLM Model: {results['config']['llm_model']}")
print(f" Knowledge Base: {results['config']['kb_type']}")
print(f"\n{'='*40} AGENTIC MODE {'='*40}")
self._print_mode_summary(results["agentic"]["metrics"])
print(f"\n{'='*40} NON-AGENTIC MODE {'='*40}")
self._print_mode_summary(results["non_agentic"]["metrics"])
print(f"\n{'='*40} COMPARISON {'='*40}")
comparison = results["comparison"]
print(f"Success Rate Difference: {comparison.get('success_rate_diff', 0):.2%} (Agentic better)")
if "time_ratio" in comparison:
print(f"Time Ratio: {comparison['time_ratio']:.2f}x (Agentic/Non-Agentic)")
print(f"Time Difference: {comparison['time_difference']:.2f} seconds")
if "keyword_recall_improvement" in comparison:
print(f"Keyword Recall Improvement: {comparison['keyword_recall_improvement']:.2%}")
if "analysis_recall_improvement" in comparison:
print(f"Analysis Recall Improvement: {comparison['analysis_recall_improvement']:.2%}")
print("\nDifficulty-Specific Improvements:")
for difficulty in ["easy", "medium", "hard"]:
key = f"{difficulty}_success_improvement"
if key in comparison:
print(f" {difficulty.capitalize()}: {comparison[key]:.2%}")
print("=" * 80)
def _print_mode_summary(self, metrics: Dict[str, Any]):
"""Print summary for a single mode"""
print(f"Total Cases: {metrics['total_cases']}")
print(f"Successful: {metrics['successful_cases']} ({metrics['successful_cases']/metrics['total_cases']:.1%})")
print(f"Failed: {metrics['failed_cases']}")
if "average_time" in metrics:
print(f"Average Time: {metrics['average_time']:.2f} seconds")
print(f"Total Time: {metrics['total_time']:.2f} seconds")
if "cases_with_citations" in metrics:
print(f"Cases with Citations: {metrics['cases_with_citations']} ({metrics['cases_with_citations']/metrics['successful_cases']:.1%})")
if "average_keyword_recall" in metrics:
print(f"Average Keyword Recall: {metrics['average_keyword_recall']:.2%}")
if "average_analysis_recall" in metrics:
print(f"Average Analysis Recall: {metrics['average_analysis_recall']:.2%}")
# Difficulty breakdown
print("\nBy Difficulty:")
for difficulty in ["easy", "medium", "hard"]:
success_key = f"{difficulty}_success_rate"
time_key = f"{difficulty}_average_time"
if success_key in metrics:
print(f" {difficulty.capitalize()}: {metrics[success_key]:.1%} success", end="")
if time_key in metrics:
print(f", {metrics[time_key]:.2f}s avg", end="")
print()
def main():
"""Main evaluation function"""
import argparse
parser = argparse.ArgumentParser(description="Evaluate Agentic RAG System")
parser.add_argument("--dataset", type=str, default="legal_qa_dataset.json",
help="Path to evaluation dataset")
parser.add_argument("--output", type=str, default="results",
help="Output directory for results")
parser.add_argument("--provider", type=str, help="Override LLM provider")
parser.add_argument("--model", type=str, help="Override LLM model")
parser.add_argument("--kb-type", choices=["local", "dify"], help="Knowledge base type")
args = parser.parse_args()
# Configure
config = Config.from_env()
if args.provider:
config.llm.provider = args.provider
if args.model:
config.llm.model = args.model
if args.kb_type:
from config import KnowledgeBaseType
config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
# Run evaluation
evaluator = RAGEvaluator(config)
results = evaluator.run_evaluation(args.dataset, args.output)
return results
if __name__ == "__main__":
main()
index_local_laws_contextual.py¶
"""Script to chunk and index local legal documents using Contextual Retrieval
This script:
1. Cleans up existing indexes
2. Reads legal documents from local laws directory
3. Chunks them with paragraph-aware boundaries (soft limit 1024, hard limit 2048)
4. Generates contextual descriptions for each chunk
5. Indexes them with contextual enhancement
"""
import os
import sys
import json
import logging
from typing import List, Dict, Any, Optional
from pathlib import Path
from dataclasses import dataclass
import time
import requests
# Add parent directory to path for imports
sys.path.append(str(Path(__file__).parent))
from config import Config, ChunkingConfig, LLMConfig
from contextual_chunking import ContextualChunker, ContextualChunk
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Configuration
RETRIEVAL_PIPELINE_URL = "http://localhost:4242" # Default retrieval pipeline URL
LAWS_DIR = Path("laws") # Local laws directory
# Custom chunking configuration for legal documents
LEGAL_CHUNKING_CONFIG = ChunkingConfig(
chunk_size=1024, # Soft limit
max_chunk_size=2048, # Hard limit
chunk_overlap=200,
respect_paragraph_boundary=True,
min_chunk_size=500
)
class ContextualLegalIndexer:
"""Handles contextual chunking and indexing of local legal documents"""
def __init__(self,
laws_dir: Path = LAWS_DIR,
pipeline_url: str = RETRIEVAL_PIPELINE_URL,
use_contextual: bool = True,
llm_config: Optional[LLMConfig] = None):
self.laws_dir = laws_dir
self.pipeline_url = pipeline_url
self.use_contextual = use_contextual
# Initialize contextual chunker
self.chunker = ContextualChunker(
chunking_config=LEGAL_CHUNKING_CONFIG,
llm_config=llm_config or LLMConfig(),
use_contextual=use_contextual
)
self.stats = {
"documents_processed": 0,
"chunks_created": 0,
"contextual_chunks": 0,
"chunks_indexed": 0,
"total_context_tokens": 0,
"errors": 0,
"categories_processed": set()
}
# Document store for tracking (must match tools.py expectation)
self.doc_store_path = Path("document_store.json")
logger.info(f"Initialized contextual indexer for local laws in {laws_dir}")
logger.info(f"Pipeline URL: {pipeline_url}")
logger.info(f"Contextual mode: {use_contextual}")
def cleanup_existing_index(self):
"""Clean up existing indexes and document store"""
logger.info("Cleaning up existing indexes...")
# Clean local document store
if self.doc_store_path.exists():
try:
# Load existing store to get document IDs
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
existing_docs = json.load(f)
logger.info(f"Found {len(existing_docs)} existing documents in store")
# Clear the store
self.doc_store_path.unlink()
logger.info("Cleared local document store")
except Exception as e:
logger.error(f"Error cleaning document store: {e}")
# Try to clear the retrieval pipeline
try:
response = requests.delete(f"{self.pipeline_url}/clear")
if response.status_code == 200:
logger.info("Cleared retrieval pipeline index")
else:
logger.warning(f"Failed to clear pipeline: {response.status_code}")
except Exception as e:
logger.warning(f"Could not clear retrieval pipeline: {e}")
logger.info("Cleanup complete")
def get_all_legal_documents(self) -> List[Dict[str, Any]]:
"""Get all legal documents from local laws directory"""
documents = []
if not self.laws_dir.exists():
logger.error(f"Laws directory not found: {self.laws_dir}")
return documents
# Iterate through category directories
for category_dir in sorted(self.laws_dir.iterdir()):
if not category_dir.is_dir():
continue
category_name = category_dir.name
logger.info(f"Found category: {category_name}")
# Find all .md files in this category
for md_file in category_dir.glob("*.md"):
doc_info = {
"path": md_file,
"name": md_file.stem, # filename without extension
"category": category_name,
"full_name": md_file.name
}
documents.append(doc_info)
logger.info(f"Found {len(documents)} legal documents")
return documents
def read_document(self, doc_info: Dict[str, Any]) -> Optional[str]:
"""Read a legal document from disk"""
try:
doc_path = doc_info["path"]
content = doc_path.read_text(encoding='utf-8')
logger.debug(f"Read {doc_info['name']} ({len(content)} chars)")
return content
except Exception as e:
logger.error(f"Error reading {doc_info['name']}: {e}")
self.stats["errors"] += 1
return None
def generate_document_id(self, doc_info: Dict[str, Any]) -> str:
"""
Generate a semantically meaningful document ID from file name.
Examples:
- "宪法.md" -> "宪法"
- "劳动法（2018-12-29）.md" -> "劳动法_2018-12-29"
- "民法典/第一编_总则.md" -> "民法典_第一编_总则"
"""
# Use the file stem (name without extension)
base_name = doc_info["name"]
# Clean up the name: remove problematic characters but keep Chinese and alphanumeric
import re
# Replace parentheses with underscore
clean_name = base_name.replace('（', '_').replace('）', '')
clean_name = clean_name.replace('(', '_').replace(')', '')
# Replace spaces and other separators with underscore
clean_name = re.sub(r'[\s\-]+', '_', clean_name)
# Remove trailing underscores
clean_name = clean_name.strip('_')
# Don't add category prefix - just use the clean file name
# This makes the IDs cleaner and more readable
doc_id = clean_name
# Ensure the ID is not too long
if len(doc_id) > 100:
# Truncate but keep it meaningful
doc_id = doc_id[:97] + "..."
logger.debug(f"Generated document ID: {doc_info['name']} -> {doc_id}")
return doc_id
def process_document(self, doc_info: Dict[str, Any], content: str,
index_immediately: bool = True) -> List[ContextualChunk]:
"""Process a document using contextual chunking with optional immediate indexing"""
# Generate semantically meaningful document ID
doc_id = self.generate_document_id(doc_info)
# Prepare document metadata
doc_metadata = {
"category": doc_info["category"],
"file_name": doc_info["full_name"],
"document_type": "legal",
"language": "zh-CN"
}
# Create callback for immediate indexing if enabled
indexed_chunks = []
def on_chunk_ready(chunk: ContextualChunk):
"""Callback to index chunk immediately when ready"""
# Update metadata
chunk.metadata["category"] = doc_info["category"]
chunk.metadata["doc_title"] = doc_info["name"]
if index_immediately:
# Index the chunk immediately
success = self.index_chunk(chunk)
if success:
logger.info(f" → Indexed chunk {chunk.chunk_index + 1} immediately")
indexed_chunks.append(chunk)
else:
logger.warning(f" → Failed to index chunk {chunk.chunk_index + 1}")
# Use the contextual chunker with callback
logger.info(f"Chunking {doc_info['name']} with contextual enhancement...")
chunks = self.chunker.chunk_document(
text=content,
doc_id=doc_id,
doc_metadata=doc_metadata,
on_chunk_ready=on_chunk_ready if index_immediately else None
)
# If not indexing immediately, update metadata for all chunks
if not index_immediately:
for chunk in chunks:
chunk.metadata["category"] = doc_info["category"]
chunk.metadata["doc_title"] = doc_info["name"]
return chunks
def index_chunk(self, chunk: ContextualChunk) -> bool:
"""Index a contextual chunk in the retrieval pipeline"""
try:
# Prepare the indexing request
# Use contextualized text for better retrieval
index_data = {
"text": chunk.contextualized_text if self.use_contextual else chunk.text,
"doc_id": chunk.chunk_id,
"metadata": {
**chunk.metadata,
"original_text": chunk.text,
"context": chunk.context,
"chunk_index": chunk.chunk_index,
"char_count": chunk.char_count,
"contextual": self.use_contextual
}
}
# Send to retrieval pipeline
response = requests.post(
f"{self.pipeline_url}/index",
json=index_data,
headers={"Content-Type": "application/json"}
)
if response.status_code == 200:
self.stats["chunks_indexed"] += 1
if chunk.context:
self.stats["contextual_chunks"] += 1
# IMPORTANT: Save to document store immediately after successful indexing
self.save_chunk_to_document_store(chunk)
return True
else:
logger.warning(f"Failed to index chunk {chunk.chunk_id}: {response.status_code}")
return False
except Exception as e:
logger.error(f"Error indexing chunk {chunk.chunk_id}: {e}")
self.stats["errors"] += 1
return False
def save_chunk_to_document_store(self, chunk: ContextualChunk):
"""Save individual chunk to document store immediately (compatible with tools.py)"""
try:
# Load existing store or create new
if self.doc_store_path.exists():
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
doc_store = json.load(f)
else:
doc_store = {}
# Save chunk in format expected by tools.py get_document method
doc_store[chunk.chunk_id] = {
"doc_id": chunk.chunk_id,
"content": chunk.contextualized_text if self.use_contextual else chunk.text,
"metadata": {
**chunk.metadata,
"original_text": chunk.text,
"context": chunk.context,
"chunk_index": chunk.chunk_index,
"char_count": chunk.char_count,
"contextual": self.use_contextual
}
}
# Write immediately - this is critical per user requirement
with open(self.doc_store_path, 'w', encoding='utf-8') as f:
json.dump(doc_store, f, ensure_ascii=False, indent=2)
logger.debug(f"Saved chunk {chunk.chunk_id} to document store")
except Exception as e:
logger.error(f"Error saving chunk {chunk.chunk_id} to document store: {e}")
def save_document_info(self, doc_info: Dict[str, Any], chunks: List[ContextualChunk]):
"""Save document summary information to local store"""
# Load existing store or create new
if self.doc_store_path.exists():
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
doc_store = json.load(f)
else:
doc_store = {}
# Add document info with meaningful ID
doc_id = self.generate_document_id(doc_info)
# Calculate statistics
total_context_tokens = sum(c.context_tokens for c in chunks)
avg_context_tokens = total_context_tokens / len(chunks) if chunks else 0
# Store document-level summary (compatible with tools.py format)
doc_store[doc_id] = {
"doc_id": doc_id,
"content": f"Document: {doc_info['name']}\nCategory: {doc_info['category']}\n\nThis is a summary entry for the complete document. Individual chunks are stored separately.",
"metadata": {
"title": doc_info["name"],
"category": doc_info["category"],
"file": str(doc_info["path"]),
"chunks": len(chunks),
"contextual_chunks": sum(1 for c in chunks if c.context),
"total_chars": sum(c.char_count for c in chunks),
"total_context_tokens": total_context_tokens,
"avg_context_tokens": avg_context_tokens,
"indexed_at": time.strftime("%Y-%m-%d %H:%M:%S"),
"chunk_ids": [c.chunk_id for c in chunks],
"is_summary": True
}
}
# Save store
with open(self.doc_store_path, 'w', encoding='utf-8') as f:
json.dump(doc_store, f, ensure_ascii=False, indent=2)
def process_all_documents(self,
max_docs: Optional[int] = None,
categories: Optional[List[str]] = None,
batch_size: int = 10):
"""Process all legal documents with contextual chunking"""
start_time = time.time()
# Clean up first
self.cleanup_existing_index()
# Get all documents
all_documents = self.get_all_legal_documents()
# Filter by categories if specified
if categories:
all_documents = [d for d in all_documents if any(cat in d["category"] for cat in categories)]
# Limit documents if specified
if max_docs:
all_documents = all_documents[:max_docs]
logger.info(f"Processing {len(all_documents)} documents...")
for i, doc_info in enumerate(all_documents):
logger.info(f"\n[{i+1}/{len(all_documents)}] Processing: {doc_info['name']}")
logger.info(f" Category: {doc_info['category']}")
# Track category
self.stats["categories_processed"].add(doc_info["category"])
# Read document
content = self.read_document(doc_info)
if not content:
continue
# Process with contextual chunking and immediate indexing
try:
# Process and index chunks immediately as they're generated
chunks = self.process_document(doc_info, content, index_immediately=True)
self.stats["chunks_created"] += len(chunks)
# Log contextual stats
if self.use_contextual:
contextual_count = sum(1 for c in chunks if c.context)
logger.info(f" ✓ Created and indexed {len(chunks)} chunks ({contextual_count} with context)")
else:
logger.info(f" ✓ Created and indexed {len(chunks)} chunks (no context)")
# Save document info immediately after processing
self.save_document_info(doc_info, chunks)
self.stats["documents_processed"] += 1
# Update token statistics
if self.use_contextual:
self.stats["total_context_tokens"] += sum(c.context_tokens for c in chunks)
except Exception as e:
logger.error(f"Error processing {doc_info['name']}: {e}")
self.stats["errors"] += 1
elapsed = time.time() - start_time
# Get final chunker statistics
chunker_stats = self.chunker.get_statistics()
# Print statistics
self._print_statistics(elapsed, chunker_stats)
def _print_statistics(self, elapsed_time: float, chunker_stats: Dict[str, Any]):
"""Print processing statistics"""
print("\n" + "="*60)
print("CONTEXTUAL INDEXING COMPLETE")
print("="*60)
print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
print(f"Time elapsed: {elapsed_time:.2f} seconds")
print(f"Categories processed: {len(self.stats['categories_processed'])}")
# Show categories
categories_list = sorted(self.stats['categories_processed'])
for cat in categories_list:
print(f" • {cat}")
print(f"\nDocuments processed: {self.stats['documents_processed']}")
print(f"Chunks created: {self.stats['chunks_created']}")
if self.use_contextual:
print(f"Contextual chunks: {self.stats['contextual_chunks']}")
print(f"Chunks indexed: {self.stats['chunks_indexed']}")
# Token usage
print(f"\nContext Generation Statistics:")
print(f" Total context tokens: {chunker_stats.get('total_context_tokens', 0):,}")
print(f" Average tokens per chunk: {chunker_stats.get('avg_context_tokens', 0):.1f}")
print(f" Total generation time: {chunker_stats.get('total_generation_time', 0):.2f}s")
print(f" Average time per chunk: {chunker_stats.get('avg_generation_time', 0):.2f}s")
# Cache statistics
print(f"\nCache Performance:")
print(f" Cache hits: {chunker_stats.get('cache_hits', 0)}")
print(f" Cache misses: {chunker_stats.get('cache_misses', 0)}")
print(f" Hit rate: {chunker_stats.get('cache_hit_rate', 0):.1%}")
# Cost estimation
print(f"\nCost Estimation:")
print(f" Estimated cost: ${chunker_stats.get('estimated_cost', 0):.2f}")
else:
print(f"Chunks indexed: {self.stats['chunks_indexed']}")
print(f"\nErrors: {self.stats['errors']}")
if self.stats['chunks_created'] > 0:
avg_chunks = self.stats['chunks_created'] / max(1, self.stats['documents_processed'])
print(f"Average chunks per document: {avg_chunks:.1f}")
if elapsed_time > 0 and self.stats['documents_processed'] > 0:
docs_per_min = (self.stats['documents_processed'] / elapsed_time) * 60
print(f"Processing speed: {docs_per_min:.1f} docs/minute")
print("="*60 + "\n")
def compare_retrieval_methods(self, test_queries: Optional[List[str]] = None):
"""Compare contextual vs non-contextual retrieval"""
if not test_queries:
test_queries = [
"什么是合同的成立条件",
"劳动者的基本权利有哪些",
"如何处理交通事故责任",
"什么是正当防卫",
"公司股东的权利和义务"
]
print("\n" + "="*60)
print("RETRIEVAL COMPARISON TEST")
print("="*60)
print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
for query in test_queries:
try:
response = requests.post(
f"{self.pipeline_url}/search",
json={
"query": query,
"mode": "hybrid",
"top_k": 10,
"rerank_top_k": 5
}
)
if response.status_code == 200:
results = response.json()
print(f"\n✓ Query: '{query}'")
if "results" in results:
print(f" Found {len(results['results'])} results")
# Show top 3 results
for i, result in enumerate(results['results'][:3]):
score = result.get('score', result.get('rerank_score', 'N/A'))
metadata = result.get('metadata', {})
print(f"\n Result {i+1}:")
print(f" Score: {score}")
print(f" Category: {metadata.get('category', 'Unknown')}")
print(f" Document: {metadata.get('doc_title', 'Unknown')}")
# Show context if available
if metadata.get('context'):
print(f" Context: {metadata['context'][:100]}...")
# Show text preview
text = result.get('text', '')
# If contextual, try to extract original text
if metadata.get('contextual') and metadata.get('original_text'):
text = metadata['original_text']
print(f" Preview: {text[:150]}...")
else:
print(f" No results found")
else:
print(f"✗ Query '{query}' failed: {response.status_code}")
except Exception as e:
print(f"✗ Error testing '{query}': {e}")
print("\n" + "="*60 + "\n")
def main():
"""Main function"""
import argparse
parser = argparse.ArgumentParser(description="Index local legal documents with contextual retrieval")
parser.add_argument("--pipeline-url", default=RETRIEVAL_PIPELINE_URL, help="Retrieval pipeline URL")
parser.add_argument("--max-docs", type=int, help="Maximum number of documents to process")
parser.add_argument("--categories", nargs="+", help="Specific categories to process")
parser.add_argument("--no-contextual", action="store_true", help="Disable contextual enhancement")
parser.add_argument("--no-cleanup", action="store_true", help="Don't clean existing indexes")
parser.add_argument("--compare", action="store_true", help="Run comparison test after indexing")
parser.add_argument("--batch-size", type=int, default=10, help="Batch size for indexing")
parser.add_argument("--llm-provider", default="kimi", help="LLM provider for context generation")
parser.add_argument("--llm-model", help="Specific LLM model to use")
args = parser.parse_args()
# Configure LLM if specified
llm_config = None
if args.llm_provider or args.llm_model:
llm_config = LLMConfig(
provider=args.llm_provider,
model=args.llm_model
)
# Create indexer
indexer = ContextualLegalIndexer(
pipeline_url=args.pipeline_url,
use_contextual=not args.no_contextual,
llm_config=llm_config
)
# Skip cleanup if requested
if args.no_cleanup:
indexer.cleanup_existing_index = lambda: logger.info("Skipping cleanup (--no-cleanup flag)")
# Process documents
indexer.process_all_documents(
max_docs=args.max_docs,
categories=args.categories,
batch_size=args.batch_size
)
# Run comparison test if requested
if args.compare:
indexer.compare_retrieval_methods()
if __name__ == "__main__":
main()
main.py¶
"""Main entry point for Agentic RAG system"""
import os
import json
import logging
import argparse
from typing import Optional
from config import Config, KnowledgeBaseType
from agent import AgenticRAG
from chunking import DocumentIndexer
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
def setup_environment():
"""Setup environment and check requirements"""
# Check for required API keys
config = Config.from_env()
# Check LLM API key
try:
api_key = config.llm.get_api_key(config.llm.provider)
if not api_key:
logger.warning(f"No API key found for provider {config.llm.provider}")
logger.info("Please set the appropriate environment variable:")
logger.info(" - MOONSHOT_API_KEY for Kimi")
logger.info(" - ARK_API_KEY for Doubao")
logger.info(" - SILICONFLOW_API_KEY for SiliconFlow")
logger.info(" - OPENAI_API_KEY for OpenAI")
return False
except Exception as e:
logger.error(f"Error checking API keys: {e}")
return False
# Check knowledge base setup
if config.knowledge_base.type == KnowledgeBaseType.LOCAL:
# Check if local retrieval pipeline is running
import requests
try:
response = requests.get(f"{config.knowledge_base.local_base_url}/health")
if response.status_code != 200:
logger.warning("Local retrieval pipeline not responding")
logger.info(f"Please ensure the retrieval pipeline is running at {config.knowledge_base.local_base_url}")
logger.info("Run: cd ../retrieval-pipeline && python main.py")
except:
logger.warning("Cannot connect to local retrieval pipeline")
logger.info("Will continue anyway - searches may fail")
elif config.knowledge_base.type == KnowledgeBaseType.DIFY:
if not config.knowledge_base.dify_api_key:
logger.warning("Dify API key not set")
logger.info("Please set DIFY_API_KEY environment variable")
return True
def run_interactive_mode(agent: AgenticRAG, mode: str = "agentic"):
"""Run interactive query mode"""
print(f"\n{'='*60}")
print(f"Agentic RAG System - {mode.capitalize()} Mode")
print(f"Verbose: {'Enabled' if agent.config.agent.verbose else 'Disabled'} | Top-K: {agent.config.knowledge_base.local_top_k}")
print(f"{'='*60}")
print("Type 'quit' or 'exit' to stop")
print("Type 'clear' to clear conversation history")
print("Type 'mode' to switch between agentic/non-agentic modes")
print(f"{'='*60}\n")
current_mode = mode
while True:
try:
user_input = input("\n[USER] > ").strip()
if user_input.lower() in ['quit', 'exit']:
print("\nGoodbye!")
break
if user_input.lower() == 'clear':
agent.clear_history()
print("Conversation history cleared.")
continue
if user_input.lower() == 'mode':
current_mode = "non-agentic" if current_mode == "agentic" else "agentic"
print(f"Switched to {current_mode} mode")
continue
if not user_input:
continue
# Process query
print(f"\n[ASSISTANT ({current_mode})] > ", end="", flush=True)
if current_mode == "agentic":
response = agent.query(user_input, stream=True)
else:
response = agent.query_non_agentic(user_input, stream=True)
# Handle streaming response
if hasattr(response, '__iter__'):
for chunk in response:
print(chunk, end="", flush=True)
print() # New line after response
else:
print(response)
except KeyboardInterrupt:
print("\n\nInterrupted. Type 'quit' to exit.")
except Exception as e:
logger.error(f"Error: {e}")
print(f"\nError processing query: {e}")
def run_batch_mode(agent: AgenticRAG, queries_file: str, output_file: str, mode: str = "agentic"):
"""Run batch queries from file"""
try:
with open(queries_file, 'r', encoding='utf-8') as f:
queries = [line.strip() for line in f if line.strip()]
except Exception as e:
logger.error(f"Error reading queries file: {e}")
return
results = []
for i, query in enumerate(queries, 1):
print(f"\n[{i}/{len(queries)}] Processing: {query[:100]}...")
try:
if mode == "agentic":
response = agent.query(query, stream=False)
else:
response = agent.query_non_agentic(query, stream=False)
results.append({
"query": query,
"response": response,
"mode": mode
})
except Exception as e:
logger.error(f"Error processing query: {e}")
results.append({
"query": query,
"response": f"Error: {str(e)}",
"mode": mode
})
# Save results
try:
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(results, f, ensure_ascii=False, indent=2)
print(f"\nResults saved to {output_file}")
except Exception as e:
logger.error(f"Error saving results: {e}")
def run_comparison_mode(agent: AgenticRAG, query: str):
"""Run both modes and compare results"""
print(f"\n{'='*60}")
print("Comparison Mode - Running both Agentic and Non-Agentic")
print(f"{'='*60}")
print(f"Query: {query}")
print(f"{'='*60}")
# Run non-agentic mode
print("\n[NON-AGENTIC MODE]")
print("-" * 40)
non_agentic_response = agent.query_non_agentic(query, stream=False)
print(non_agentic_response)
# Clear history for fair comparison
agent.clear_history()
# Run agentic mode
print("\n[AGENTIC MODE]")
print("-" * 40)
agentic_response = agent.query(query, stream=False)
print(agentic_response)
print(f"\n{'='*60}")
def main():
"""Main function"""
parser = argparse.ArgumentParser(description="Agentic RAG System")
# Mode selection
parser.add_argument("--mode", choices=["agentic", "non-agentic", "compare"],
default="agentic", help="Query mode")
# Query options
parser.add_argument("--query", type=str, help="Single query to process")
parser.add_argument("--batch", type=str, help="Path to file with queries (one per line)")
parser.add_argument("--output", type=str, default="results.json",
help="Output file for batch results")
# Configuration options
parser.add_argument("--provider", type=str, help="LLM provider")
parser.add_argument("--model", type=str, help="LLM model")
parser.add_argument("--kb-type", choices=["local", "dify"], help="Knowledge base type")
parser.add_argument("--verbose", action="store_true", help="Verbose output (default: True)")
parser.add_argument("--no-verbose", action="store_true", help="Disable verbose output")
# Indexing options
parser.add_argument("--index", type=str, help="Path to index (file or directory)")
parser.add_argument("--chunk-size", type=int, default=2048, help="Chunk size for indexing")
args = parser.parse_args()
# Setup environment
if not setup_environment():
logger.warning("Environment setup incomplete, continuing anyway...")
# Load or create config
config = Config.from_env()
# Set verbose mode by default (can be disabled with --no-verbose)
config.agent.verbose = True # Default to verbose mode
# Override config with command line args
if args.provider:
config.llm.provider = args.provider
if args.model:
config.llm.model = args.model
if args.kb_type:
config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
# Handle verbose mode (default is True, can be disabled with --no-verbose)
if args.no_verbose:
config.agent.verbose = False
elif args.verbose:
config.agent.verbose = True # Explicitly set if --verbose is used
# Handle indexing if requested
if args.index:
print(f"\n{'='*60}")
print("Indexing Documents")
print(f"{'='*60}")
config.chunking.chunk_size = args.chunk_size
indexer = DocumentIndexer(config.knowledge_base, config.chunking)
from pathlib import Path
path = Path(args.index)
if path.is_file():
result = indexer.index_file(str(path))
elif path.is_dir():
result = indexer.index_directory(str(path))
else:
print(f"Path not found: {path}")
return
print(json.dumps(result, indent=2, ensure_ascii=False))
print(f"{'='*60}\n")
# Create agent
agent = AgenticRAG(config)
# Handle different execution modes
if args.query and args.mode == "compare":
# Comparison mode with single query
run_comparison_mode(agent, args.query)
elif args.query:
# Single query mode
print(f"\n[Query] {args.query}")
print(f"[Mode] {args.mode}")
print(f"[Verbose] {'Enabled' if config.agent.verbose else 'Disabled'}")
print(f"[Top-K] {config.knowledge_base.local_top_k}")
print("-" * 40)
if args.mode == "agentic":
response = agent.query(args.query, stream=False)
else:
response = agent.query_non_agentic(args.query, stream=False)
print(response)
elif args.batch:
# Batch mode
run_batch_mode(agent, args.batch, args.output, args.mode)
else:
# Interactive mode (default)
run_interactive_mode(agent, args.mode)
if __name__ == "__main__":
main()
quickstart.py¶
#!/usr/bin/env python3
"""Quick start script to test the Contextual Retrieval System
This script provides a quick way to test contextual retrieval
with a sample document and see the improvements.
"""
import logging
from pathlib import Path
from config import Config
from contextual_chunking import ContextualChunker
from contextual_tools import ContextualKnowledgeBaseTools
# Simple logging for quickstart
logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger(__name__)
def main():
"""Quick demonstration of contextual retrieval"""
print("\n" + "="*60)
print("CONTEXTUAL RETRIEVAL - QUICK START")
print("="*60 + "\n")
# Sample document about multiple companies
document = """
2023 Technology Sector Report
Apple Inc. Performance:
Apple reported exceptional results in 2023. The company's revenue reached
$394 billion, with iPhone sales contributing 52% of total revenue. The
services division showed strong growth of 16% year-over-year. Tim Cook
emphasized the company's commitment to innovation and sustainability.
Microsoft Corporation Update:
Microsoft achieved record cloud revenue in 2023. Azure revenue grew by 27%
as enterprises accelerated digital transformation. The company's total
revenue was $211 billion. CEO Satya Nadella highlighted AI integration
across all product lines as a key strategic priority.
Google (Alphabet) Highlights:
Google's parent company Alphabet reported $283 billion in revenue for 2023.
Search advertising remained the largest revenue driver at $175 billion.
The company increased AI research spending by 30% to maintain competitive
advantage. YouTube advertising revenue exceeded $40 billion.
Market Analysis:
The technology sector showed resilience despite economic headwinds. Companies
that invested heavily in AI and cloud infrastructure outperformed the market.
The sector's average growth rate was 12%, with cloud services growing at 25%
and traditional hardware declining by 3%.
"""
print("Step 1: Initializing systems...")
config = Config.from_env()
# Create both contextual and non-contextual systems
contextual_chunker = ContextualChunker(use_contextual=True)
non_contextual_chunker = ContextualChunker(use_contextual=False)
contextual_kb = ContextualKnowledgeBaseTools(use_contextual=True)
non_contextual_kb = ContextualKnowledgeBaseTools(use_contextual=False)
print("\nStep 2: Processing document...")
print("-" * 40)
# Process with contextual system
print("Creating contextual chunks...")
contextual_chunks = contextual_chunker.chunk_document(
text=document,
doc_id="tech_report_2023"
)
contextual_kb.index_contextual_chunks(contextual_chunks)
print(f"✓ Created {len(contextual_chunks)} contextual chunks")
# Process with non-contextual system
print("Creating non-contextual chunks...")
non_contextual_chunks = non_contextual_chunker.chunk_document(
text=document,
doc_id="tech_report_2023"
)
non_contextual_kb.index_contextual_chunks(non_contextual_chunks)
print(f"✓ Created {len(non_contextual_chunks)} non-contextual chunks")
# Show example contextual chunk
if contextual_chunks:
print("\nExample Contextual Chunk:")
print("-" * 40)
chunk = contextual_chunks[0]
print(f"Original text: {chunk.text[:100]}...")
print(f"Added context: {chunk.context}")
print("\n" + "="*60)
print("Step 3: Testing Search Queries")
print("="*60)
# Test queries
queries = [
"What was the company's revenue?",
"Which company emphasized AI?",
"What was the growth rate?"
]
for query in queries:
print(f"\nQuery: '{query}'")
print("-" * 40)
# Contextual search
contextual_results = contextual_kb.contextual_search(query, top_k=1)
# Non-contextual search
non_contextual_results = non_contextual_kb.contextual_search(query, top_k=1)
print("\nContextual Result:")
if contextual_results:
result = contextual_results[0]
print(f" Score: {result.score:.4f}")
if result.context_text:
print(f" Context: {result.context_text[:80]}...")
print(f" Match: {result.text[:100]}...")
else:
print(" No results")
print("\nNon-Contextual Result:")
if non_contextual_results:
result = non_contextual_results[0]
print(f" Score: {result.score:.4f}")
print(f" Match: {result.text[:100]}...")
else:
print(" No results")
# Compare scores
if contextual_results and non_contextual_results:
improvement = ((contextual_results[0].score - non_contextual_results[0].score)
/ non_contextual_results[0].score * 100)
print(f"\n📊 Improvement: {improvement:+.1f}%")
print("\n" + "="*60)
print("SUMMARY")
print("="*60)
# Get statistics
stats = contextual_chunker.get_statistics()
print(f"\nContextual Chunking Statistics:")
print(f" Chunks processed: {stats['total_chunks']}")
print(f" Context tokens used: {stats['total_context_tokens']}")
print(f" Estimated cost: ${stats['estimated_cost']:.4f}")
print("\nKey Insights:")
print("✓ Contextual chunks preserve company-specific information")
print("✓ Ambiguous queries ('the company') are resolved correctly")
print("✓ Search accuracy improves significantly with context")
print("\n" + "="*60)
print("Quick start complete! Try with your own documents:")
print(" python contextual_main.py --mode index --document your_file.txt")
print("="*60 + "\n")
if __name__ == "__main__":
main()
test_document_ids.py¶
test_simple.py¶
#!/usr/bin/env python3
"""Simple test script for Agentic RAG system"""
import os
import json
from pathlib import Path
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
def test_basic_functionality():
"""Test basic functionality of the system"""
print("🧪 Testing Agentic RAG System")
print("="*60)
# Import modules
try:
from config import Config, KnowledgeBaseType
from agent import AgenticRAG
from tools import KnowledgeBaseTools
from chunking import DocumentChunker, DocumentIndexer
print("✅ All modules imported successfully")
except ImportError as e:
print(f"❌ Import error: {e}")
return False
# Test configuration
print("\n📋 Testing Configuration...")
try:
config = Config.from_env()
print(f" Provider: {config.llm.provider}")
print(f" KB Type: {config.knowledge_base.type}")
print(f" Chunk Size: {config.chunking.chunk_size}")
print("✅ Configuration loaded")
except Exception as e:
print(f"❌ Config error: {e}")
return False
# Test document chunking
print("\n📄 Testing Document Chunking...")
try:
chunker = DocumentChunker(config.chunking)
sample_text = """故意杀人罪是指故意非法剥夺他人生命的行为。
根据《中华人民共和国刑法》第二百三十二条规定，故意杀人的，
处死刑、无期徒刑或者十年以上有期徒刑；情节较轻的，
处三年以上十年以下有期徒刑。
量刑考虑因素包括犯罪动机、手段、后果等。"""
chunks = chunker.chunk_text(sample_text, "test_doc")
print(f" Created {len(chunks)} chunks")
print(f" First chunk: {chunks[0]['text'][:100]}...")
print("✅ Chunking works")
except Exception as e:
print(f"❌ Chunking error: {e}")
return False
# Test knowledge base tools
print("\n🔧 Testing Knowledge Base Tools...")
try:
kb_tools = KnowledgeBaseTools(config.knowledge_base)
# Add test document to store
kb_tools.add_document(
"test_doc_1",
"故意杀人罪处死刑、无期徒刑或者十年以上有期徒刑。",
{"source": "test"}
)
# Test document retrieval
doc = kb_tools.get_document("test_doc_1")
if "error" not in doc:
print(f" Retrieved document: {doc['doc_id']}")
print("✅ Document storage works")
else:
print(f"⚠️ Document retrieval returned: {doc}")
except Exception as e:
print(f"❌ KB Tools error: {e}")
return False
# Test agent initialization
print("\n🤖 Testing Agent Initialization...")
try:
agent = AgenticRAG(config)
print(f" Model: {agent.model}")
print(f" Provider: {config.llm.provider}")
print("✅ Agent initialized")
except Exception as e:
print(f"❌ Agent initialization error: {e}")
print(" Make sure you have set the appropriate API key in .env")
return False
# Test simple query (if API key is available)
if os.getenv("MOONSHOT_API_KEY") or os.getenv("OPENAI_API_KEY"):
print("\n💬 Testing Simple Query...")
try:
# Add some test data
kb_tools.add_document(
"criminal_law_test",
"""盗窃罪的立案标准：
1. 数额较大：一般为1000元至3000元以上
2. 多次盗窃：2年内盗窃3次以上
3. 入户盗窃、携带凶器盗窃、扒窃不论数额""",
{"type": "law"}
)
# Test non-agentic query (simpler, less likely to fail)
response = agent.query_non_agentic("盗窃罪立案标准", stream=False)
if response and len(response) > 10:
print(f" Response: {response[:200]}...")
print("✅ Query processing works")
else:
print(f"⚠️ Response was empty or too short: {response}")
except Exception as e:
print(f"⚠️ Query error: {e}")
print(" This might be due to retrieval pipeline not running")
else:
print("\n⚠️ Skipping query test (no API key found)")
print("\n" + "="*60)
print("🎉 Basic functionality test complete!")
return True
def test_evaluation_dataset():
"""Test evaluation dataset generation"""
print("\n📊 Testing Evaluation Dataset...")
try:
# Import dataset builder
import sys
sys.path.append("evaluation")
from dataset_builder import LegalDatasetBuilder, create_legal_documents
# Build dataset
builder = LegalDatasetBuilder()
simple_cases = builder.create_simple_cases()
complex_cases = builder.create_complex_cases()
print(f" Simple cases: {len(simple_cases)}")
print(f" Complex cases: {len(complex_cases)}")
print(f" First simple case: {simple_cases[0]['question']}")
# Create documents
documents = create_legal_documents()
print(f" Legal documents: {len(documents)}")
print("✅ Evaluation dataset works")
return True
except Exception as e:
print(f"❌ Dataset error: {e}")
return False
if __name__ == "__main__":
print("🚀 Agentic RAG System - Test Suite")
print("="*60)
# Run tests
success = test_basic_functionality()
if success:
test_evaluation_dataset()
print("\n" + "="*60)
if success:
print("✅ All basic tests passed!")
print("\nNext steps:")
print("1. Make sure retrieval pipeline is running:")
print(" cd ../retrieval-pipeline && python main.py")
print("\n2. Run the quickstart:")
print(" python quickstart.py")
print("\n3. Or start interactive mode:")
print(" python main.py")
else:
print("❌ Some tests failed. Please check the errors above.")
print("\nCommon issues:")
print("1. Missing API keys in .env file")
print("2. Retrieval pipeline not running")
print("3. Missing dependencies (run: pip install -r requirements.txt)")
tools.py¶
"""Tools for knowledge base interaction"""
import json
import logging
import requests
from typing import Dict, Any, List, Optional
from dataclasses import dataclass
from config import KnowledgeBaseConfig, KnowledgeBaseType
logger = logging.getLogger(__name__)
@dataclass
class SearchResult:
"""Search result from knowledge base"""
doc_id: str
chunk_id: str
text: str
score: float
metadata: Dict[str, Any] = None
def to_dict(self) -> Dict[str, Any]:
return {
"doc_id": self.doc_id,
"chunk_id": self.chunk_id,
"text": self.text,
"score": self.score,
"metadata": self.metadata or {}
}
class KnowledgeBaseTools:
"""Tools for interacting with knowledge base"""
def __init__(self, config: KnowledgeBaseConfig):
self.config = config
self.document_store = {} # In-memory store for documents
# Load document store if exists
try:
with open(config.document_store_path, 'r', encoding='utf-8') as f:
self.document_store = json.load(f)
except FileNotFoundError:
logger.info("No existing document store found, starting fresh")
except Exception as e:
logger.error(f"Error loading document store: {e}")
def save_document_store(self):
"""Save document store to disk"""
try:
with open(self.config.document_store_path, 'w', encoding='utf-8') as f:
json.dump(self.document_store, f, ensure_ascii=False, indent=2)
except Exception as e:
logger.error(f"Error saving document store: {e}")
def knowledge_base_search(self, query: str) -> List[Dict[str, Any]]:
"""
Search the knowledge base with a natural language query.
Args:
query: Natural language query string
Returns:
List of matching document chunks with scores
"""
try:
logger.info(f"Knowledge base search initiated - Type: {self.config.type}, Query: '{query}'")
if self.config.type == KnowledgeBaseType.LOCAL:
return self._search_local(query)
elif self.config.type == KnowledgeBaseType.DIFY:
return self._search_dify(query)
elif self.config.type == KnowledgeBaseType.RAPTOR:
return self._search_raptor(query)
elif self.config.type == KnowledgeBaseType.GRAPHRAG:
return self._search_graphrag(query)
else:
logger.error(f"Unsupported knowledge base type: {self.config.type}")
raise ValueError(f"Unsupported knowledge base type: {self.config.type}")
except ValueError:
raise # Re-raise ValueError for unsupported types
except Exception as e:
logger.error(f"Error in knowledge base search: {e}")
return []
def _search_local(self, query: str) -> List[Dict[str, Any]]:
"""Search using local retrieval pipeline"""
try:
logger.info(f"Searching local knowledge base for: {query}")
response = requests.post(
f"{self.config.local_base_url}/search",
json={
"query": query,
"mode": "hybrid",
"top_k": self.config.local_top_k,
"rerank": True
},
timeout=30 # Add timeout
)
response.raise_for_status()
results = []
data = response.json()
# The retrieval pipeline returns results in 'reranked_results' field for hybrid mode
results_field = data.get("reranked_results") or data.get("results") or []
if not results_field:
# Check other possible fields
if data.get("dense_results"):
results_field = data.get("dense_results", [])
elif data.get("sparse_results"):
results_field = data.get("sparse_results", [])
if not results_field:
logger.warning(f"Search returned empty results for query: {query}")
logger.debug(f"Response keys: {data.keys()}")
return []
for item in results_field:
# Extract doc_id - the field name varies between results
doc_id = item.get("doc_id", "")
# Use doc_id as chunk_id since the retrieval pipeline indexes chunks
chunk_id = doc_id
# Get text from the result
text = item.get("text", "")
# Get score - might be 'rerank_score' or 'score'
score = item.get("rerank_score", item.get("score", 0.0))
result = SearchResult(
doc_id=doc_id,
chunk_id=chunk_id,
text=text,
score=score,
metadata=item.get("metadata", {})
)
results.append(result.to_dict())
logger.info(f"Local search returned {len(results)} results for query: {query}")
return results
except requests.exceptions.Timeout as e:
logger.error(f"Timeout connecting to local retrieval pipeline at {self.config.local_base_url}: {e}")
return []
except requests.exceptions.ConnectionError as e:
logger.error(f"Cannot connect to local retrieval pipeline at {self.config.local_base_url}: {e}")
logger.info("Make sure the retrieval pipeline is running: cd ../retrieval-pipeline && python main.py")
return []
except requests.exceptions.RequestException as e:
logger.error(f"Error connecting to local retrieval pipeline: {e}")
return []
def _search_dify(self, query: str) -> List[Dict[str, Any]]:
"""Search using Dify API"""
if not self.config.dify_api_key:
logger.error("Dify API key not configured")
return []
try:
headers = {
"Authorization": f"Bearer {self.config.dify_api_key}",
"Content-Type": "application/json"
}
payload = {
"query": query,
"top_k": self.config.dify_top_k
}
if self.config.dify_dataset_id:
payload["dataset_id"] = self.config.dify_dataset_id
response = requests.post(
f"{self.config.dify_base_url}/datasets/search",
headers=headers,
json=payload
)
response.raise_for_status()
results = []
data = response.json()
for item in data.get("data", {}).get("records", []):
doc_id = item.get("document_id", "")
chunk_id = item.get("segment_id", f"{doc_id}_chunk_{len(results)}")
result = SearchResult(
doc_id=doc_id,
chunk_id=chunk_id,
text=item.get("content", ""),
score=item.get("score", 0.0),
metadata=item.get("metadata", {})
)
results.append(result.to_dict())
logger.info(f"Dify search returned {len(results)} results")
return results
except requests.exceptions.RequestException as e:
logger.error(f"Error connecting to Dify API: {e}")
return []
def _search_raptor(self, query: str) -> List[Dict[str, Any]]:
"""Search using RAPTOR tree-based index"""
try:
response = requests.post(
f"{self.config.raptor_base_url}/query",
json={
"query": query,
"index_type": "raptor",
"top_k": self.config.raptor_top_k
}
)
response.raise_for_status()
results = []
data = response.json()
for i, item in enumerate(data.get("results", [])):
# RAPTOR returns tree nodes with levels and summaries
doc_id = item.get("node_id", f"raptor_node_{i}")
chunk_id = f"{doc_id}_level_{item.get('level', 0)}"
# Use summary if available, otherwise use text
text_content = item.get("summary", item.get("text", ""))
result = SearchResult(
doc_id=doc_id,
chunk_id=chunk_id,
text=text_content,
score=item.get("score", 0.0),
metadata={
"level": item.get("level", 0),
"source": "raptor"
}
)
results.append(result.to_dict())
logger.info(f"RAPTOR search returned {len(results)} results")
return results
except requests.exceptions.RequestException as e:
logger.error(f"Error connecting to RAPTOR index: {e}")
return []
def _search_graphrag(self, query: str) -> List[Dict[str, Any]]:
"""Search using GraphRAG knowledge graph index"""
try:
response = requests.post(
f"{self.config.graphrag_base_url}/query",
json={
"query": query,
"index_type": "graphrag",
"top_k": self.config.graphrag_top_k,
"search_type": self.config.graphrag_search_type
}
)
response.raise_for_status()
results = []
data = response.json()
for i, item in enumerate(data.get("results", [])):
# GraphRAG returns entities or communities
result_type = item.get("type", "unknown")
if result_type == "entity":
doc_id = item.get("id", f"entity_{i}")
chunk_id = f"{doc_id}_{item.get('entity_type', 'unknown')}"
text_content = f"{item.get('name', '')}. {item.get('description', '')}"
metadata = {
"type": "entity",
"entity_type": item.get("entity_type"),
"related_entities": item.get("related_entities", [])
}
else: # community
doc_id = item.get("id", f"community_{i}")
chunk_id = f"{doc_id}_level_{item.get('level', 0)}"
text_content = item.get("summary", "")
metadata = {
"type": "community",
"level": item.get("level", 0),
"entity_count": item.get("entity_count", 0),
"sample_entities": item.get("sample_entities", [])
}
result = SearchResult(
doc_id=doc_id,
chunk_id=chunk_id,
text=text_content,
score=item.get("score", 0.0),
metadata={**metadata, "source": "graphrag"}
)
results.append(result.to_dict())
logger.info(f"GraphRAG search returned {len(results)} results")
return results
except requests.exceptions.RequestException as e:
logger.error(f"Error connecting to GraphRAG index: {e}")
return []
def get_document(self, doc_id: str) -> Dict[str, Any]:
"""
Retrieve the entire document from the knowledge base.
Args:
doc_id: Document ID
Returns:
Full document content and metadata
"""
try:
# First check local document store
if doc_id in self.document_store:
return self.document_store[doc_id]
if self.config.type == KnowledgeBaseType.LOCAL:
return self._get_document_local(doc_id)
elif self.config.type == KnowledgeBaseType.DIFY:
return self._get_document_dify(doc_id)
elif self.config.type == KnowledgeBaseType.RAPTOR:
return self._get_document_raptor(doc_id)
elif self.config.type == KnowledgeBaseType.GRAPHRAG:
return self._get_document_graphrag(doc_id)
else:
raise ValueError(f"Unsupported knowledge base type: {self.config.type}")
except Exception as e:
logger.error(f"Error retrieving document {doc_id}: {e}")
return {"error": f"Document {doc_id} not found"}
def _get_document_local(self, doc_id: str) -> Dict[str, Any]:
"""Get document from local retrieval pipeline"""
try:
response = requests.get(
f"{self.config.local_base_url}/documents/{doc_id}"
)
if response.status_code == 404:
return {"error": f"Document {doc_id} not found"}
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
logger.error(f"Error getting document from local pipeline: {e}")
return {"error": str(e)}
def _get_document_dify(self, doc_id: str) -> Dict[str, Any]:
"""Get document from Dify"""
if not self.config.dify_api_key:
return {"error": "Dify API key not configured"}
try:
headers = {
"Authorization": f"Bearer {self.config.dify_api_key}",
"Content-Type": "application/json"
}
response = requests.get(
f"{self.config.dify_base_url}/documents/{doc_id}",
headers=headers
)
if response.status_code == 404:
return {"error": f"Document {doc_id} not found"}
response.raise_for_status()
return response.json()
except requests.exceptions.RequestException as e:
logger.error(f"Error getting document from Dify: {e}")
return {"error": str(e)}
def _get_document_raptor(self, doc_id: str) -> Dict[str, Any]:
"""Get document/node from RAPTOR index"""
try:
# For RAPTOR, we perform a targeted search for the specific node
response = requests.post(
f"{self.config.raptor_base_url}/query",
json={
"query": f"node:{doc_id}", # Specific node query
"index_type": "raptor",
"top_k": 1
}
)
if response.status_code == 404:
return {"error": f"Document {doc_id} not found"}
response.raise_for_status()
data = response.json()
if data.get("results"):
result = data["results"][0]
return {
"doc_id": doc_id,
"content": result.get("text", ""),
"metadata": {
"summary": result.get("summary", ""),
"level": result.get("level", 0),
"source": "raptor"
}
}
return {"error": f"Document {doc_id} not found"}
except requests.exceptions.RequestException as e:
logger.error(f"Error getting document from RAPTOR: {e}")
return {"error": str(e)}
def _get_document_graphrag(self, doc_id: str) -> Dict[str, Any]:
"""Get entity or community from GraphRAG index"""
try:
# For GraphRAG, we perform a targeted search for the specific entity/community
response = requests.post(
f"{self.config.graphrag_base_url}/query",
json={
"query": f"id:{doc_id}", # Specific ID query
"index_type": "graphrag",
"top_k": 1,
"search_type": "hybrid"
}
)
if response.status_code == 404:
return {"error": f"Document {doc_id} not found"}
response.raise_for_status()
data = response.json()
if data.get("results"):
result = data["results"][0]
content = ""
metadata = {"source": "graphrag"}
if result.get("type") == "entity":
content = f"{result.get('name', '')}\n\n{result.get('description', '')}"
metadata.update({
"type": "entity",
"entity_type": result.get("entity_type"),
"related_entities": result.get("related_entities", [])
})
else: # community
content = result.get("summary", "")
metadata.update({
"type": "community",
"level": result.get("level", 0),
"entity_count": result.get("entity_count", 0),
"sample_entities": result.get("sample_entities", [])
})
return {
"doc_id": doc_id,
"content": content,
"metadata": metadata
}
return {"error": f"Document {doc_id} not found"}
except requests.exceptions.RequestException as e:
logger.error(f"Error getting document from GraphRAG: {e}")
return {"error": str(e)}
def add_document(self, doc_id: str, content: str, metadata: Optional[Dict] = None):
"""Add a document to the local store"""
self.document_store[doc_id] = {
"doc_id": doc_id,
"content": content,
"metadata": metadata or {}
}
self.save_document_store()
# Tool function definitions for agent
def get_tool_definitions() -> List[Dict[str, Any]]:
"""Get OpenAI-format tool definitions"""
return [
{
"type": "function",
"function": {
"name": "knowledge_base_search",
"description": "Search the knowledge base for relevant information using a natural language query. Returns top-matching document chunks.",
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Natural language search query to find relevant information"
}
},
"required": ["query"]
}
}
},
{
"type": "function",
"function": {
"name": "get_document",
"description": "Retrieve the complete content of a specific document from the knowledge base using its document ID.",
"parameters": {
"type": "object",
"properties": {
"doc_id": {
"type": "string",
"description": "The unique identifier of the document to retrieve"
}
},
"required": ["doc_id"]
}
}
}
]
evaluation/retrieval_eval.json¶
{
"description": "上下文感知检索 vs. 传统分块的检索召回评测集（对应实验 3-11）。gold_chunk_id 为人工标注的相关文本块，标注依据见 note 字段。语料为 document_store.json 中已建好索引的《宪法》与《检察官法》分块。",
"corpus": "document_store.json",
"queries": [
{
"id": "q01",
"query": "宪法是哪一年通过的？后来经过了几次修正？",
"gold_chunk_id": "宪法_chunk_0",
"note": "chunk_0 为宪法开篇，列举 1982 年通过日期及历次修正案日期。"
},
{
"id": "q02",
"query": "国家的根本任务是什么？",
"gold_chunk_id": "宪法_chunk_1",
"note": "chunk_1 序言明确写明“国家的根本任务是……”。"
},
{
"id": "q03",
"query": "我国处理民族关系的基本原则和对外政策中的五项原则是什么？",
"gold_chunk_id": "宪法_chunk_2",
"note": "chunk_2 阐述社会主义民族关系与和平共处五项原则。"
},
{
"id": "q04",
"query": "少数民族聚居的地方如何实行区域自治？",
"gold_chunk_id": "宪法_chunk_3",
"note": "chunk_3 第四条前后规定民族区域自治与依法治国。"
},
{
"id": "q05",
"query": "宪法如何保护个体经济、私营经济等非公有制经济？",
"gold_chunk_id": "宪法_chunk_4",
"note": "chunk_4 含第十一至十三条，规定非公有制经济与私有财产保护。"
},
{
"id": "q06",
"query": "宪法关于设立特别行政区以及在华外国人权利的规定是什么？",
"gold_chunk_id": "宪法_chunk_6",
"note": "chunk_6 含第三十一、三十二条，特别行政区与外国人条款。"
},
{
"id": "q07",
"query": "公民对国家机关工作人员的违法失职行为享有哪些监督权利？",
"gold_chunk_id": "宪法_chunk_7",
"note": "chunk_7 第四十一条规定批评、建议、申诉、控告、检举权。"
},
{
"id": "q08",
"query": "全国人民代表大会常务委员会有哪些职权？",
"gold_chunk_id": "宪法_chunk_10",
"note": "旗舰对照案例：chunk_10 原文以“（四）解释法律……”开头，正文既无“常务委员会”也无“职权”字样，纯靠上下文前缀（“第六十七条 全国人民代表大会常务委员会行使下列职权”）才能被锚定；无上下文时极易被 chunk_9/chunk_11 抢占。"
},
{
"id": "q09",
"query": "国家主席有哪些职权？",
"gold_chunk_id": "宪法_chunk_12",
"note": "chunk_12 含第八十至八十二条，国家主席公布法律、任免国务院总理等职权。"
},
{
"id": "q10",
"query": "地方各级人民代表大会的任期和代表选举方式是怎样规定的？",
"gold_chunk_id": "宪法_chunk_14",
"note": "chunk_14 第九十七至九十九条规定地方人大代表选举与五年任期。"
},
{
"id": "q11",
"query": "县级以上地方各级人民代表大会常务委员会组成人员有什么兼职限制？",
"gold_chunk_id": "宪法_chunk_15",
"note": "chunk_15 规定常委会组成人员不得担任行政、监察、审判、检察机关职务。"
},
{
"id": "q12",
"query": "监察委员会的性质、组成和任期是怎样规定的？",
"gold_chunk_id": "宪法_chunk_17",
"note": "chunk_17 第七节监察委员会，规定其性质、组成与任期。"
},
{
"id": "q13",
"query": "检察官法是什么时候修订的？",
"gold_chunk_id": "检察官法_2019_04_23_chunk_0",
"note": "chunk_0 列举检察官法 1995 年通过及历次修正、2019 年修订日期。"
},
{
"id": "q14",
"query": "检察官遴选委员会由哪些人员组成？律师参加公开选拔需要什么条件？",
"gold_chunk_id": "检察官法_2019_04_23_chunk_2",
"note": "chunk_2 第十六条规定省级检察官遴选委员会组成及律师公开选拔条件。"
},
{
"id": "q15",
"query": "检察官在哪些情形下应当予以免职？",
"gold_chunk_id": "检察官法_2019_04_23_chunk_3",
"note": "chunk_3 列举检察官免职情形及违反条件任命的撤销程序。"
}
]
}