agentic-rag¶
第3章 · 用户记忆和知识库 · 配套项目
chapter3/agentic-rag
项目说明¶
Agentic RAG System¶
An educational implementation of an Agentic Retrieval-Augmented Generation (RAG) system with ReAct pattern, supporting multiple LLM providers and knowledge base backends.
🌟 Features¶
- Agentic RAG with ReAct Pattern: Uses reasoning and tool-calling to iteratively search and retrieve information
- Non-Agentic RAG Mode: Simple retrieval + LLM response for comparison
- Multiple LLM Provider Support:
- Kimi/Moonshot
- Doubao
- SiliconFlow
- OpenAI
- OpenRouter
- Groq
- Together AI
- DeepSeek
- Flexible Knowledge Base:
- Offline BM25 (内置，零依赖离线运行): in-process BM25 over the bundled
laws/corpus — no server, no API key required - Local retrieval pipeline (requires ../retrieval-pipeline)
- Dify knowledge base API
- Document Chunking: Configurable chunking with paragraph boundary respect
- Evaluation Framework: Comprehensive evaluation with Chinese legal dataset
- Conversation History: Support for follow-up questions
- Verbose Logging: Detailed logging to understand agent reasoning
📦 Installation¶
⚙️ Configuration¶
Set environment variables in .env file:
# LLM API Keys (set the one you're using)
MOONSHOT_API_KEY=your_kimi_api_key
ARK_API_KEY=your_doubao_api_key
SILICONFLOW_API_KEY=your_siliconflow_api_key
OPENAI_API_KEY=your_openai_api_key
OPENROUTER_API_KEY=your_openrouter_api_key
GROQ_API_KEY=your_groq_api_key
TOGETHER_API_KEY=your_together_api_key
DEEPSEEK_API_KEY=your_deepseek_api_key
# Knowledge Base Configuration (optional, defaults to local)
KB_TYPE=local # Options: "offline" (内置离线 BM25，无需服务/API), "local", "dify"
DIFY_API_KEY=your_dify_api_key # if using Dify
DIFY_DATASET_ID=your_dataset_id # optional
# LLM Configuration (optional)
LLM_PROVIDER=kimi # default provider
LLM_MODEL=kimi-k3 # optional, uses provider defaults
🚀 Usage¶
0. 零依赖离线对比实验（推荐先跑，无需 API / 无需外部服务）¶
本实验的核心论点是：面对复杂问题，让 Agent 自主分解、多轮迭代检索，其证据召回显著优于单次检索。
compare_offline.py 用内置的离线 BM25 检索器（offline_retriever.py，直接读取 laws/ 语料）
在小型中文司法问答集上量化这一差距，完全离线、无需任何 API Key：
python compare_offline.py
# 可选参数：--corpus laws --top-k 5 --dataset evaluation/offline_qa.json --output result.json
真实输出（本机实测，21372 个法条分块 / 288 篇文档）：
问题 难度 单次检索 分解检索 检索次数
------------------------------------------------------------------------------
故意伤害致人重伤的，如何处… easy 100% 100% 1 → 1
正当防卫是怎么规定的？ easy 100% 100% 1 → 1
醉酒驾驶机动车如何处罚？ easy 100% 100% 1 → 1
故意杀人罪判几年？ hard 0% 100% 1 → 1
盗窃罪的立案标准是什么？ hard 0% 100% 1 → 1
诈骗罪的量刑标准是什么？ hard 0% 100% 1 → 1
醉酒过失致人重伤且有盗窃前… hard 33% 100% 1 → 3
------------------------------------------------------------------------------
聚合指标（平均证据召回率）:
全部 48% 100% 1.0 → 1.3
简单题 100% 100% 1.0 → 1.0
复杂题 8% 100% 1.0 → 1.5
解读（与书中实验 3-9 一致）：简单问题两种范式相差无几（均 100%），一次直接检索就够；
复杂/措辞欠佳的问题上差距显著（8% → 100%），单次检索因关键词不精确而漏检关键法条，
分解式多轮检索则能逐一补齐证据。该指标为纯检索层的『证据召回率』，是回答质量的上界
——检索不到证据，生成阶段无从谈起。金标准法条均已确认存在于 laws/ 语料中。
说明：离线模式用数据集中预先标注的
subqueries表示『Agent 分解后发起的检索』， 以隔离出检索策略本身的贡献；在真实系统中，这些子查询由 LLM 在 ReAct 循环中动态生成。 需要 LLM 生成、端到端评测答案质量时，请使用evaluation/evaluate.py（需配置 API Key）。
也可以让完整 Agent 直接跑在离线知识库上（检索离线，仅答案生成需要 API Key）：
python main.py --kb-type offline --query "醉酒过失致人重伤且有盗窃前科如何量刑"
python main.py --kb-type offline --query "故意杀人罪判几年" --mode compare
1. Start the Retrieval Pipeline¶
First, start the retrieval pipeline server (required for local knowledge base):
# In a separate terminal
cd ../retrieval-pipeline
python main.py
# Server will run on http://localhost:4242
2. Index Documents¶
Option A: Index Chinese Law Documents (Pre-included)¶
# Index the included Chinese law documents
python index_local_laws.py
# With specific categories
python index_local_laws.py --categories 宪法 民法典
# With document limit
python index_local_laws.py --max-docs 10
Option B: Index Custom Documents¶
# Index a single file
python main.py --index path/to/document.txt
# Index a directory
python main.py --index path/to/documents/
# Custom chunk size
python main.py --index documents/ --chunk-size 2048
3. Run the Agentic RAG System¶
Interactive Mode (Default)¶
# Start in agentic mode (default)
python main.py
# Start in non-agentic mode
python main.py --mode non-agentic
# Enable verbose logging (default is enabled)
python main.py --verbose
# Disable verbose logging
python main.py --no-verbose
In interactive mode: - Type your questions and press Enter - Type 'quit' or 'exit' to stop - Type 'clear' to clear conversation history - Type 'mode' to switch between agentic/non-agentic modes
Single Query¶
# Agentic mode
python main.py --query "宪法第一条是什么？" --mode agentic
# Non-agentic mode
python main.py --query "盗窃罪的立案标准是什么？" --mode non-agentic
# Compare both modes
python main.py --query "故意杀人罪判几年？" --mode compare
Batch Processing¶
# Create a file with queries (one per line)
echo "故意杀人罪判几年？
盗窃罪的立案标准是什么？
醉酒驾驶如何处罚？" > queries.txt
# Run batch
python main.py --batch queries.txt --output results.json
# Batch with specific mode
python main.py --batch queries.txt --mode non-agentic
With Different Providers¶
python main.py --provider openai --model gpt-5.6-luna
python main.py --provider doubao --model doubao-seed-1-6-thinking-250715
python main.py --provider siliconflow --query "你好"
4. Run Evaluation¶
# Build evaluation dataset
cd evaluation
python dataset_builder.py
# Run evaluation
python evaluate.py
# With specific configuration
python evaluate.py --provider kimi --kb-type local --output custom_results
📁 Project Structure¶
agentic-rag/
├── config.py # Configuration classes
├── agent.py # Main AgenticRAG implementation
├── tools.py # Knowledge base tools (含 offline BM25 后端)
├── offline_retriever.py # 内置离线 BM25 检索器（读取 laws/，无需服务/API）
├── compare_offline.py # 离线对比实验：分解检索 vs 单次检索（证据召回率表）
├── chunking.py # Document chunking and indexing
├── main.py # Main entry point
├── index_local_laws.py # Index Chinese law documents
├── quickstart.py # Quick demo script
├── test_simple.py # Simple test script
├── requirements.txt # Dependencies
├── README.md # This file
├── document_store.json # Local document storage
├── laws/ # Chinese law documents
│ ├── 1-宪法/
│ ├── 2-宪法相关法/
│ ├── 3-民法典/
│ ├── 3-民法商法/
│ ├── 4-行政法/
│ ├── 5-经济法/
│ ├── 6-社会法/
│ ├── 7-刑法/
│ └── 8-诉讼与非诉讼程序法/
└── evaluation/
├── dataset_builder.py # Build evaluation dataset
├── offline_qa.json # 离线对比数据集（问题 + 金标准法条 + Agent 分解子查询）
└── evaluate.py # Evaluation framework (端到端答案质量，需 API)
🧠 How It Works¶
Agentic RAG Mode¶
The agent uses the ReAct (Reasoning + Acting) pattern:
- Reasoning: Analyzes what information is needed to answer the question
- Tool Calling: Uses
knowledge_base_searchtool to find relevant chunks - Iterative Search: May perform multiple searches with refined queries
- Document Retrieval: Can fetch complete documents with
get_documentfor context - Answer Synthesis: Combines retrieved information with citations
- Conversation Memory: Maintains context for follow-up questions
Example flow:
User: 宪法第一条是什么？
Agent: [Thinks] Need to find information about Article 1 of the Constitution
[Tool] knowledge_base_search("宪法第一条")
[Result] Found relevant chunks about constitutional articles
[Answer] Based on the retrieved information, Article 1 states...
Non-Agentic RAG Mode¶
Simple retrieval-augmented generation:
- Direct Search: Searches once with the user's query as-is
- Context Injection: Puts top-K results in the prompt
- Single Response: LLM answers based on provided context
- No Iteration: Single-shot approach without refinement
📊 Configuration Options¶
Top-K Results¶
Control how many search results to retrieve:
Verbose Mode¶
See detailed agent reasoning:
# Enable verbose (default)
python main.py --verbose
# Disable for cleaner output
python main.py --no-verbose
LLM Temperature¶
Control response randomness:
🎯 Evaluation Results¶
检索层（离线、可复现、真实实测）：见上文 第 0 节 的证据召回率表——分解式多轮检索把复杂题的召回率从 8% 提升到 100%，而简单题两种范式打平（均 100%）。
生成层（端到端答案质量，需 LLM API）：evaluation/evaluate.py 在此基础上真正调用 LLM 生成答案，
统计关键词/分析点召回、引用覆盖率、响应时间等指标。以下为该框架产出的指标与典型模式：
Metrics¶
- Success Rate: Whether the answer contains key legal concepts
- Response Time: Time to generate response
- Retrieval Quality: Relevance of retrieved chunks
- Citation Coverage: Proper source attribution
Expected Patterns¶
Agentic RAG typically shows:
- ✅ Better coverage of complex multi-faceted questions
- ✅ More accurate citations through explicit tool use
- ✅ Ability to refine searches based on initial results
- ⚠️ Slower response time (multiple LLM calls)
Non-Agentic RAG typically shows: - ✅ Faster responses (single retrieval step) - ✅ Good performance on simple, direct questions - ⚠️ May miss relevant information with poor query formulation - ⚠️ Limited ability to handle ambiguous queries
🔧 Troubleshooting¶
Retrieval Pipeline Not Responding¶
# Check if the service is running
curl http://localhost:4242/health
# If not, start it:
cd ../retrieval-pipeline
python main.py
No Search Results¶
-
Ensure documents are indexed:
-
Check document store:
-
Verify retrieval pipeline has documents:
API Key Issues¶
# Check if environment variable is set
echo $MOONSHOT_API_KEY
# Or use .env file
cat .env | grep API_KEY
Indexing Errors¶
- Ensure retrieval pipeline is running before indexing
- Check file encodings (UTF-8 expected)
- Verify network connectivity to localhost:4242
🤝 Contributing¶
Areas for potential enhancement:
- Additional evaluation metrics
- More sophisticated chunking strategies
- Better reranking algorithms
- Additional knowledge base backends (RAPTOR, GraphRAG)
- Multi-language support
- Query expansion techniques
- Hybrid retrieval strategies
📄 License¶
This is an educational project for learning purposes.
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 _is_reasoning_model(model) -> bool:
"""Whether the model is a reasoning model (Kimi K3, GPT-5, ...)."""
m = str(model or "").lower().replace("/", "-")
return "kimi-k3" in m or "gpt-5" in m
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."""
return 1 if _is_reasoning_model(model) else requested
def _reasoning_safe_max_tokens(model, requested=1024, floor=4096):
"""Reasoning models spend part of their budget on hidden reasoning tokens,
so a small ``max_tokens`` (e.g. 1024) silently truncates the visible answer.
Ensure reasoning models get at least ``floor`` tokens; leave other providers
at the requested value."""
if _is_reasoning_model(model):
return max(requested, floor)
return 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", "")
results = self.kb_tools.knowledge_base_search(query)
# Log full trajectory when verbose
if self.config.agent.verbose:
logger.info("=" * 80)
logger.info(f"TOOL EXECUTION: {tool_name}")
logger.info("-" * 80)
logger.info(f"Query: {query}")
logger.info("-" * 80)
if not results:
if self.config.agent.verbose:
logger.info("Results: No relevant documents found")
logger.info("=" * 80)
return {"status": "no_results", "message": "No relevant documents found"}
# Format results for agent - KEEP ALL RESULTS
formatted_results = []
for i, r in enumerate(results, 1):
formatted_results.append({
"doc_id": r["doc_id"],
"chunk_id": r["chunk_id"],
"text": r["text"],
"score": r["score"]
})
# Log each result in full detail
if self.config.agent.verbose:
logger.info(f"Result {i}/{len(results)}:")
logger.info(f" Document ID: {r['doc_id']}")
logger.info(f" Chunk ID: {r['chunk_id']}")
logger.info(f" Score: {r['score']:.4f}")
logger.info(f" Text (full):\n{'-' * 40}")
logger.info(r['text'])
logger.info("-" * 40)
if self.config.agent.verbose:
logger.info(f"Total results found: {len(results)}")
logger.info("=" * 80)
return {
"status": "success",
"results": formatted_results[:3], # Limit to top 3 for LLM context
"total_found": len(results),
"all_results": formatted_results # Keep all for logging
}
elif tool_name == "get_document":
doc_id = arguments.get("doc_id", "")
# Log full trajectory when verbose
if self.config.agent.verbose:
logger.info("=" * 80)
logger.info(f"TOOL EXECUTION: {tool_name}")
logger.info("-" * 80)
logger.info(f"Document ID: {doc_id}")
logger.info("-" * 80)
document = self.kb_tools.get_document(doc_id)
if "error" in document:
if self.config.agent.verbose:
logger.info(f"Error: {document['error']}")
logger.info("=" * 80)
return {"status": "error", "message": document["error"]}
# Log full document content
if self.config.agent.verbose:
logger.info("Document Retrieved:")
logger.info(f" Doc ID: {document.get('doc_id', doc_id)}")
if document.get('metadata'):
logger.info(f" Metadata: {json.dumps(document['metadata'], indent=2, ensure_ascii=False)}")
logger.info(" Content (full):\n" + "=" * 40)
logger.info(document.get('content', ''))
logger.info("=" * 80)
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("\n" + "=" * 100)
logger.info(f"ITERATION {iterations}/{max_iterations}")
logger.info("=" * 100)
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=_reasoning_safe_max_tokens(self.model, 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:
arguments = {}
if self.config.agent.verbose:
logger.info("\n" + "#" * 80)
logger.info(f"TOOL CALL: {tool_name}")
logger.info(f"Arguments: {json.dumps(arguments, indent=2, ensure_ascii=False)}")
logger.info("#" * 80)
# Execute tool
result = self._execute_tool(tool_name, arguments)
# Log full tool result when verbose
if self.config.agent.verbose:
logger.info("\n" + "*" * 80)
logger.info("TOOL RESULT:")
logger.info("*" * 80)
# Show full result including all_results if present
if 'all_results' in result:
logger.info("All Search Results (Complete):")
for idx, res in enumerate(result['all_results'], 1):
logger.info(f"\nResult {idx}:")
logger.info(json.dumps(res, indent=2, ensure_ascii=False))
else:
logger.info(json.dumps(result, indent=2, ensure_ascii=False))
logger.info("*" * 80 + "\n")
# For messages, don't include all_results to avoid overloading LLM
result_for_llm = {k: v for k, v in result.items() if k != 'all_results'}
# Add tool result to messages
tool_message = {
"role": "tool",
"tool_call_id": tool_call.id,
"content": json.dumps(result_for_llm, 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)
# Return response
if stream:
return self._stream_response(message.content or "")
else:
return message.content or ""
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
search_results = self.kb_tools.knowledge_base_search(user_query)
# Build context from search results
context_parts = []
for i, result in enumerate(search_results[:3], 1): # Top 3 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=_reasoning_safe_max_tokens(self.model, 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=_reasoning_safe_max_tokens(self.model, 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_offline.py¶
"""离线对比实验：智能体化 RAG（多轮/分解检索）vs 非智能体化 RAG（单次检索）。
本脚本**完全离线运行**——只做检索、不调用任何 LLM、不依赖外部检索服务，
因此无需 API Key 即可复现。它在一个小型中文司法问答集（evaluation/offline_qa.json）
上，量化对比两种检索范式的『证据召回率』：
- 非智能体化：把用户原始问题作为唯一查询做一次检索（single-shot）；
- 智能体化：模拟 Agent 分解/改写问题后发起多次检索，再对结果取并集。
金标准（gold_articles）为回答每个问题所必需的法条编号；某法条被判定为『命中』
当且仅当检索结果中存在一个以该法条编号开头的分块。证据召回率 = 命中金标准法条数
/ 金标准法条总数。这一检索层指标是回答质量的上界：检索不到证据，生成阶段就无从
谈起。生成阶段的端到端评测（需要 LLM API）见 evaluation/evaluate.py。
"""
import os
import re
import sys
import json
import time
import argparse
from typing import List, Dict, Any
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from offline_retriever import OfflineRetriever, _ARTICLE_RE
def _leading_article(text: str) -> str:
"""抽取分块开头的法条编号（如『第二百三十五条』），无则返回空串。"""
m = _ARTICLE_RE.match(text.strip())
return m.group(0) if m else ""
def _covered(retrieved: List[Dict[str, Any]], gold_articles: List[str]) -> List[str]:
"""返回被检索结果命中的金标准法条列表。"""
hit_markers = {_leading_article(r["text"]) for r in retrieved}
hit_markers.discard("")
return [g for g in gold_articles if g in hit_markers]
def run_case(retriever: OfflineRetriever, case: Dict[str, Any], top_k: int) -> Dict[str, Any]:
gold = case["gold_articles"]
# 非智能体化：单次检索，查询即用户原始问题。
naive_query = case.get("naive_query", case["question"])
naive_hits = retriever.search(naive_query, top_k)
naive_covered = _covered(naive_hits, gold)
# 智能体化：分解为多个子查询，逐一检索后取并集。
subqueries = case.get("subqueries") or [case["question"]]
agentic_hits: List[Dict[str, Any]] = []
seen = set()
for sq in subqueries:
for r in retriever.search(sq, top_k):
if r["chunk_id"] not in seen:
seen.add(r["chunk_id"])
agentic_hits.append(r)
agentic_covered = _covered(agentic_hits, gold)
return {
"id": case["id"],
"question": case["question"],
"difficulty": case.get("difficulty", "unknown"),
"gold_articles": gold,
"naive": {
"num_searches": 1,
"covered": naive_covered,
"recall": len(naive_covered) / len(gold) if gold else 0.0,
},
"agentic": {
"num_searches": len(subqueries),
"covered": agentic_covered,
"recall": len(agentic_covered) / len(gold) if gold else 0.0,
},
}
def _mean(xs: List[float]) -> float:
return sum(xs) / len(xs) if xs else 0.0
def _pad(text: str, width: int) -> str:
"""按显示宽度左对齐（一个中文字符按两个宽度计）。"""
display = sum(2 if ord(c) > 127 else 1 for c in text)
return text + " " * max(0, width - display)
def summarize(results: List[Dict[str, Any]]) -> Dict[str, Any]:
def agg(subset):
return {
"count": len(subset),
"naive_recall": _mean([r["naive"]["recall"] for r in subset]),
"agentic_recall": _mean([r["agentic"]["recall"] for r in subset]),
"naive_searches": _mean([r["naive"]["num_searches"] for r in subset]),
"agentic_searches": _mean([r["agentic"]["num_searches"] for r in subset]),
}
summary = {"overall": agg(results)}
for diff in ("easy", "hard"):
subset = [r for r in results if r["difficulty"] == diff]
if subset:
summary[diff] = agg(subset)
return summary
def print_table(results: List[Dict[str, Any]], summary: Dict[str, Any]):
print("\n" + "=" * 78)
print("离线检索对比：证据召回率（Evidence Recall）")
print("=" * 78)
print(_pad("问题", 30) + _pad("难度", 8) + _pad("单次检索", 12)
+ _pad("分解检索", 12) + "检索次数")
print("-" * 78)
for r in results:
q = (r["question"][:13] + "…") if len(r["question"]) > 13 else r["question"]
naive = f"{r['naive']['recall']:.0%}"
agentic = f"{r['agentic']['recall']:.0%}"
searches = f"1 → {r['agentic']['num_searches']}"
print(_pad(q, 30) + _pad(r["difficulty"], 8) + _pad(naive, 12)
+ _pad(agentic, 12) + searches)
print("-" * 78)
def row(name, s):
print(_pad(name, 30) + _pad("", 8) + _pad(f"{s['naive_recall']:.0%}", 12)
+ _pad(f"{s['agentic_recall']:.0%}", 12)
+ f"{s['naive_searches']:.1f} → {s['agentic_searches']:.1f}")
print("聚合指标（平均证据召回率）:")
row(" 全部", summary["overall"])
if "easy" in summary:
row(" 简单题", summary["easy"])
if "hard" in summary:
row(" 复杂题", summary["hard"])
print("=" * 78)
ov = summary["overall"]
lift = ov["agentic_recall"] - ov["naive_recall"]
print(f"结论：分解式多轮检索将整体证据召回率从 {ov['naive_recall']:.0%} "
f"提升到 {ov['agentic_recall']:.0%}（+{lift:.0%}），"
f"代价是平均检索次数由 {ov['naive_searches']:.1f} 增至 {ov['agentic_searches']:.1f}。")
if "hard" in summary:
hv = summary["hard"]
print(f" 复杂题上的差距最为显著：{hv['naive_recall']:.0%} → {hv['agentic_recall']:.0%}。")
print("=" * 78)
def main():
parser = argparse.ArgumentParser(
description="离线对比智能体化 RAG（多轮分解检索）与非智能体化 RAG（单次检索）的证据召回率；纯检索、无需 LLM 与外部服务。",
formatter_class=argparse.RawTextHelpFormatter,
)
parser.add_argument("--dataset", type=str, default="evaluation/offline_qa.json",
help="离线问答数据集路径（默认：evaluation/offline_qa.json）")
parser.add_argument("--corpus", type=str, default="laws",
help="法律语料目录，用于构建离线 BM25 索引（默认：laws）")
parser.add_argument("--top-k", type=int, default=5,
help="每次检索返回的分块数量，即检索深度（默认：5）")
parser.add_argument("--output", type=str, default=None,
help="将详细结果写入的 JSON 文件路径（默认：不落盘，仅打印）")
args = parser.parse_args()
print(f"[离线对比] 构建 BM25 索引，语料目录：{args.corpus} …")
t0 = time.time()
retriever = OfflineRetriever(args.corpus)
print(f"[离线对比] 索引完成：{len(retriever.chunks)} 个法条分块 / "
f"{len(retriever.documents)} 篇文档，用时 {time.time() - t0:.1f}s")
with open(args.dataset, "r", encoding="utf-8") as f:
dataset = json.load(f)
cases = dataset["cases"]
results = [run_case(retriever, c, args.top_k) for c in cases]
summary = summarize(results)
print_table(results, summary)
if args.output:
payload = {
"dataset": args.dataset,
"corpus": args.corpus,
"top_k": args.top_k,
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
"results": results,
"summary": summary,
}
with open(args.output, "w", encoding="utf-8") as f:
json.dump(payload, 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"""
OFFLINE = "offline" # In-process BM25 over local law corpus (no server, no API)
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.
if not api_key and provider_lower != "openrouter" and os.getenv("OPENROUTER_API_KEY"):
model = _openrouter_model_id(self.model or defaults.get("model"))
return {
"api_key": os.getenv("OPENROUTER_API_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
# Offline in-process BM25 backend config (no external server / no API key)
offline_corpus_path: str = "laws"
offline_top_k: int = 5
# 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 = 3
# RAPTOR tree-based index config
raptor_base_url: str = "http://localhost:4242"
raptor_top_k: int = 3
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 = 3
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
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.py¶
"""Script to chunk and index local legal documents from laws directory
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. Indexes them in the retrieval pipeline
"""
import os
import json
import logging
import hashlib
from typing import List, Dict, Any, Optional
from pathlib import Path
from dataclasses import dataclass
import time
import requests
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
# Chunking configuration
SOFT_LIMIT = 1024 # Soft character limit
HARD_LIMIT = 2048 # Hard character limit
MIN_CHUNK_SIZE = 500 # Minimum chunk size
@dataclass
class LegalChunk:
"""Represents a chunk of legal document"""
chunk_id: str
doc_id: str
doc_title: str
category: str # e.g., "宪法", "民法典"
text: str
chunk_index: int
char_count: int
metadata: Dict[str, Any]
class LocalLegalIndexer:
"""Handles chunking and indexing of local legal documents"""
def __init__(self, laws_dir: Path = LAWS_DIR, pipeline_url: str = RETRIEVAL_PIPELINE_URL):
self.laws_dir = laws_dir
self.pipeline_url = pipeline_url
self.stats = {
"documents_processed": 0,
"chunks_created": 0,
"chunks_indexed": 0,
"errors": 0,
"categories_processed": set()
}
# Document store for tracking
self.doc_store_path = Path("document_store.json")
logger.info(f"Initialized indexer for local laws in {laws_dir}")
logger.info(f"Pipeline URL: {pipeline_url}")
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"Processing 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 across {len(self.stats['categories_processed'])} categories")
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 chunk_document_smart(self, text: str, doc_id: str, doc_title: str, category: str) -> List[LegalChunk]:
"""
Smart chunking that respects paragraph boundaries with soft and hard limits.
Strategy:
- Accumulate paragraphs until soft limit is exceeded
- Keep adding if next paragraph fits within hard limit
- Cut at paragraph boundary when possible
- Force split at hard limit if necessary
"""
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 > HARD_LIMIT:
# Save current chunk if exists
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
if len(chunk_text) >= MIN_CHUNK_SIZE:
chunks.append(self._create_chunk(
chunk_text, doc_id, doc_title, category, len(chunks)
))
current_chunk = []
current_size = 0
# Force split the oversized paragraph
for i in range(0, para_size, SOFT_LIMIT):
sub_text = para[i:i + SOFT_LIMIT]
if len(sub_text) >= MIN_CHUNK_SIZE:
chunks.append(self._create_chunk(
sub_text, doc_id, doc_title, category, len(chunks)
))
continue
# Check if adding this paragraph would exceed limits
new_size = current_size + para_size + (4 if current_chunk else 0) # Account for \n\n
if new_size > SOFT_LIMIT and current_chunk:
# Check if we can still fit it under hard limit
if new_size <= HARD_LIMIT:
# Add it anyway (between soft and hard limit)
current_chunk.append(para)
current_size = new_size
else:
# Save current chunk and start new one
chunk_text = '\n\n'.join(current_chunk)
if len(chunk_text) >= MIN_CHUNK_SIZE:
chunks.append(self._create_chunk(
chunk_text, doc_id, doc_title, category, len(chunks)
))
# Start new chunk
current_chunk = [para]
current_size = para_size
else:
# Add to current chunk
current_chunk.append(para)
current_size = new_size
# Save final chunk
if current_chunk:
chunk_text = '\n\n'.join(current_chunk)
if len(chunk_text) >= MIN_CHUNK_SIZE:
chunks.append(self._create_chunk(
chunk_text, doc_id, doc_title, category, len(chunks)
))
logger.info(f"Created {len(chunks)} chunks for {doc_title}")
return chunks
def _create_chunk(self, text: str, doc_id: str, doc_title: str, category: str, chunk_index: int) -> LegalChunk:
"""Create a LegalChunk object"""
chunk_id = f"{doc_id}_chunk_{chunk_index}"
# Extract section info if available
section_info = self._extract_section_info(text)
return LegalChunk(
chunk_id=chunk_id,
doc_id=doc_id,
doc_title=doc_title,
category=category,
text=text,
chunk_index=chunk_index,
char_count=len(text),
metadata={
"source": "local_laws",
"document_type": "legal",
"language": "zh-CN",
"category": category,
"section": section_info
}
)
def _extract_section_info(self, text: str) -> Optional[str]:
"""Extract section/chapter/article information from legal text"""
import re
# Common patterns in Chinese legal documents
patterns = [
r'第[一二三四五六七八九十百千\d]+[章节条款篇编]',
r'第[一二三四五六七八九十百千\d]+部分',
r'[【\[]第[一二三四五六七八九十百千\d]+[章节条款篇编][】\]]'
]
for pattern in patterns:
match = re.search(pattern, text[:200]) # Check first 200 chars
if match:
return match.group()
return None
def index_chunk(self, chunk: LegalChunk) -> bool:
"""Index a chunk in the retrieval pipeline"""
try:
# Prepare the indexing request
index_data = {
"text": chunk.text,
"doc_id": chunk.chunk_id,
"metadata": {
**chunk.metadata,
"doc_title": chunk.doc_title,
"category": chunk.category,
"chunk_index": chunk.chunk_index,
"char_count": chunk.char_count
}
}
# 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
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_document_info(self, doc_info: Dict[str, Any], chunks: List[LegalChunk]):
"""Save document 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
doc_id = hashlib.md5(doc_info["full_name"].encode()).hexdigest()[:12]
doc_store[doc_id] = {
"title": doc_info["name"],
"category": doc_info["category"],
"file": str(doc_info["path"]),
"chunks": len(chunks),
"total_chars": sum(c.char_count for c in chunks),
"indexed_at": time.strftime("%Y-%m-%d %H:%M:%S")
}
# 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):
"""Process all legal documents"""
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']}")
# Track category
self.stats["categories_processed"].add(doc_info["category"])
# Read document
content = self.read_document(doc_info)
if not content:
continue
# Generate document ID
doc_id = hashlib.md5(doc_info["full_name"].encode()).hexdigest()[:12]
# Chunk the document
chunks = self.chunk_document_smart(
content,
doc_id,
doc_info["name"],
doc_info["category"]
)
self.stats["chunks_created"] += len(chunks)
# Index each chunk
indexed_count = 0
for j, chunk in enumerate(chunks):
if self.index_chunk(chunk):
indexed_count += 1
# Progress update
if (j + 1) % 10 == 0:
logger.debug(f" Indexed {j + 1}/{len(chunks)} chunks")
logger.info(f" ✓ Indexed {indexed_count}/{len(chunks)} chunks successfully")
# Save document info
self.save_document_info(doc_info, chunks)
self.stats["documents_processed"] += 1
elapsed = time.time() - start_time
# Print statistics
self._print_statistics(elapsed)
def _print_statistics(self, elapsed_time: float):
"""Print processing statistics"""
print("\n" + "="*60)
print("INDEXING COMPLETE")
print("="*60)
print(f"Time elapsed: {elapsed_time:.2f} seconds")
print(f"Categories processed: {len(self.stats['categories_processed'])}")
print(f" - {', '.join(sorted(self.stats['categories_processed']))}")
print(f"Documents processed: {self.stats['documents_processed']}")
print(f"Chunks created: {self.stats['chunks_created']}")
print(f"Chunks indexed: {self.stats['chunks_indexed']}")
print(f"Errors: {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_sec = self.stats['documents_processed'] / elapsed_time
print(f"Processing speed: {docs_per_sec:.2f} docs/second")
print("="*60 + "\n")
def verify_indexing(self, test_queries: Optional[List[str]] = None):
"""Verify that indexing worked by performing test searches"""
if not test_queries:
test_queries = [
"民法典",
"合同法",
"劳动法",
"刑法",
"宪法"
]
print("\n" + "="*60)
print("VERIFICATION TESTS")
print("="*60)
for query in test_queries:
try:
response = requests.post(
f"{self.pipeline_url}/search",
json={
"query": query,
"mode": "hybrid",
"top_k": 5,
"rerank_top_k": 3
}
)
if response.status_code == 200:
results = response.json()
print(f"\n✓ Test search for '{query}':")
if "results" in results:
print(f" Found {len(results['results'])} results")
for i, result in enumerate(results['results'][:2]):
score = result.get('score', result.get('rerank_score', 'N/A'))
metadata = result.get('metadata', {})
print(f" {i+1}. Score: {score}")
print(f" Category: {metadata.get('category', 'Unknown')}")
print(f" Doc: {metadata.get('doc_title', 'Unknown')}")
print(f" Preview: {result.get('text', '')[:100]}...")
else:
print(f" No results found")
else:
print(f"✗ Test search for '{query}' failed: {response.status_code}")
except Exception as e:
print(f"✗ Error testing '{query}': {e}")
print("="*60 + "\n")
def main():
"""Main function"""
import argparse
parser = argparse.ArgumentParser(description="Index local legal documents into retrieval pipeline")
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 (e.g., '宪法' '民法典')")
parser.add_argument("--no-cleanup", action="store_true", help="Don't clean existing indexes")
parser.add_argument("--verify", action="store_true", help="Run verification tests after indexing")
args = parser.parse_args()
# Create indexer
indexer = LocalLegalIndexer(pipeline_url=args.pipeline_url)
# 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
)
# Run verification if requested
if args.verify:
indexer.verify_indexing()
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"""
kb = agent.config.knowledge_base
active_top_k = kb.offline_top_k if kb.type == KnowledgeBaseType.OFFLINE else kb.local_top_k
print(f"\n{'='*60}")
print(f"Agentic RAG System - {mode.capitalize()} Mode")
print(f"Verbose: {'Enabled' if agent.config.agent.verbose else 'Disabled'} | KB: {kb.type.value} | Top-K: {active_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="智能体化 RAG 系统：对比『智能体化（多轮迭代检索）』与『非智能体化（单次检索）』两种范式。",
epilog=(
"示例:\n"
" python main.py --kb-type offline --query \"醉酒过失致人重伤且有盗窃前科如何量刑\"\n"
" python main.py --query \"故意杀人罪判几年\" --mode compare --kb-type offline\n"
" python compare_offline.py # 纯离线检索对比，无需 API 与外部服务\n"
),
formatter_class=argparse.RawTextHelpFormatter,
)
# 模式选择
parser.add_argument("--mode", choices=["agentic", "non-agentic", "compare"],
default="agentic",
help="查询模式：agentic=智能体化多轮检索 / non-agentic=单次检索 / compare=同题对比（默认：agentic）")
# 查询选项
parser.add_argument("--query", type=str, help="单条查询问题；不指定则进入交互模式")
parser.add_argument("--batch", type=str, help="批量查询文件路径（每行一个问题）")
parser.add_argument("--output", type=str, default="results.json",
help="批量结果的输出文件路径（默认：results.json）")
# 配置选项
parser.add_argument("--provider", type=str, help="LLM 提供商（如 kimi / doubao / openai）")
parser.add_argument("--model", type=str, help="LLM 模型名称（不指定则用提供商默认模型）")
parser.add_argument("--kb-type", choices=["offline", "local", "dify"],
help="知识库后端：offline=内置离线 BM25（无需服务/无需 API）/ local=检索流水线服务 / dify=Dify API")
parser.add_argument("--corpus", type=str,
help="离线后端的法律语料目录（仅 --kb-type offline 生效，默认：laws）")
parser.add_argument("--top-k", type=int, dest="top_k",
help="检索深度：每次检索返回的分块数量（默认：offline=5，local=3）")
parser.add_argument("--verbose", action="store_true", help="输出详细的 Agent 推理轨迹（默认开启）")
parser.add_argument("--no-verbose", action="store_true", help="关闭详细日志输出")
# 索引选项
parser.add_argument("--index", type=str, help="待索引的文件或目录路径")
parser.add_argument("--chunk-size", type=int, default=2048, help="索引时的分块大小（字符数，默认：2048）")
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)
if args.corpus:
config.knowledge_base.offline_corpus_path = args.corpus
if args.top_k:
# 同时设置离线与本地后端的检索深度，保持行为一致
config.knowledge_base.offline_top_k = args.top_k
config.knowledge_base.local_top_k = args.top_k
# 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
kb = config.knowledge_base
active_top_k = kb.offline_top_k if kb.type == KnowledgeBaseType.OFFLINE else kb.local_top_k
print(f"\n[Query] {args.query}")
print(f"[Mode] {args.mode}")
print(f"[KB] {kb.type.value}")
print(f"[Verbose] {'Enabled' if config.agent.verbose else 'Disabled'}")
print(f"[Top-K] {active_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()
offline_retriever.py¶
"""In-process offline retriever (BM25 over the local law corpus).
This backend makes the whole experiment runnable without the external
`retrieval-pipeline` HTTP service: it reads the Markdown law files under
``laws/``, splits them into article-level chunks (每一条法条一个 chunk), and
scores queries with Okapi BM25. Retrieval therefore runs fully offline with no
API key and no server; only the LLM answer-generation step (in ``agent.py``)
still needs a provider API.
Chinese text is tokenised with ``jieba`` when available, falling back to a
character uni/bi-gram tokeniser so the module works with only the standard
library installed.
"""
import os
import re
import math
import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
from collections import Counter, defaultdict
logger = logging.getLogger(__name__)
# Article marker at the start of a line, e.g. 第二百三十五条 / 第一百三十三条之一
_ARTICLE_RE = re.compile(r"^第[一二三四五六七八九十百千零两0-9]+条(?:之[一二三四五六七八九十0-9]+)?")
def _tokenize(text: str) -> List[str]:
"""Tokenise mixed Chinese/English text.
Prefers jieba; otherwise emits ASCII words plus Chinese character uni- and
bi-grams, which is enough for lexical BM25 matching without extra deps.
"""
try:
import jieba # type: ignore
return [t for t in jieba.cut(text) if t.strip()]
except Exception:
tokens: List[str] = []
for m in re.findall(r"[a-zA-Z0-9]+|[一-鿿]+", text):
if m[0].isascii():
tokens.append(m.lower())
else:
tokens.extend(list(m)) # unigrams
tokens.extend(m[i:i + 2] for i in range(len(m) - 1)) # bigrams
return tokens
class OfflineRetriever:
"""Okapi BM25 retriever over article-level chunks of the law corpus."""
def __init__(self,
corpus_path: str = "laws",
k1: float = 1.5,
b: float = 0.75,
extensions: Optional[List[str]] = None):
self.corpus_path = corpus_path
self.k1 = k1
self.b = b
self.extensions = extensions or [".md", ".txt"]
self.chunks: List[Dict[str, Any]] = [] # {doc_id, chunk_id, title, category, text}
self.documents: Dict[str, Dict[str, Any]] = {} # doc_id -> {title, category, file, content}
self._doc_freqs: List[Counter] = [] # per-chunk term frequencies
self._doc_lens: List[int] = []
self._df: Dict[str, int] = defaultdict(int) # document frequency per term
self._idf: Dict[str, float] = {}
self._avg_len: float = 0.0
self._build_index()
# ------------------------------------------------------------------ build
def _iter_files(self):
root = Path(self.corpus_path)
if not root.exists():
logger.warning(f"Offline corpus path not found: {root}")
return
for path in sorted(root.rglob("*")):
if path.is_file() and path.suffix in self.extensions:
yield path
def _split_articles(self, content: str) -> List[str]:
"""Split a law document into article-level chunks.
Falls back to blank-line paragraph grouping when the file has no
``第X条`` markers (e.g. non-statute documents).
"""
lines = content.splitlines()
articles: List[str] = []
current: List[str] = []
seen_article = False
for line in lines:
if _ARTICLE_RE.match(line.strip()):
seen_article = True
if current:
articles.append("\n".join(current).strip())
current = [line]
else:
current.append(line)
if current:
articles.append("\n".join(current).strip())
if not seen_article:
# No article markers: group by blank lines into ~paragraph chunks.
articles = [p.strip() for p in content.split("\n\n") if p.strip()]
return [a for a in articles if a]
def _build_index(self):
for path in self._iter_files():
try:
content = path.read_text(encoding="utf-8")
except Exception as e:
logger.error(f"Error reading {path}: {e}")
continue
category = path.parent.name
title = path.stem
doc_id = f"{category}/{title}"
self.documents[doc_id] = {
"doc_id": doc_id,
"title": title,
"category": category,
"file": str(path),
"content": content,
}
for idx, article in enumerate(self._split_articles(content)):
if len(article) < 4:
continue
chunk_id = f"{doc_id}_chunk_{idx}"
self.chunks.append({
"doc_id": doc_id,
"chunk_id": chunk_id,
"title": title,
"category": category,
"text": article,
})
# Build BM25 statistics.
for chunk in self.chunks:
tf = Counter(_tokenize(chunk["text"]))
self._doc_freqs.append(tf)
self._doc_lens.append(sum(tf.values()))
for term in tf:
self._df[term] += 1
n = len(self.chunks)
self._avg_len = (sum(self._doc_lens) / n) if n else 0.0
for term, df in self._df.items():
# BM25 idf with +1 to stay non-negative.
self._idf[term] = math.log(1 + (n - df + 0.5) / (df + 0.5))
logger.info(
f"OfflineRetriever indexed {n} chunks from {len(self.documents)} "
f"documents under '{self.corpus_path}'"
)
# ----------------------------------------------------------------- search
def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
"""Return the ``top_k`` article chunks scored by BM25 for ``query``."""
if not self.chunks:
return []
q_terms = _tokenize(query)
scored: List[tuple] = []
for i, tf in enumerate(self._doc_freqs):
dl = self._doc_lens[i]
score = 0.0
for term in q_terms:
f = tf.get(term)
if not f:
continue
idf = self._idf.get(term, 0.0)
denom = f + self.k1 * (1 - self.b + self.b * dl / (self._avg_len or 1))
score += idf * (f * (self.k1 + 1)) / denom
if score > 0:
scored.append((score, i))
scored.sort(reverse=True)
results: List[Dict[str, Any]] = []
for score, i in scored[:top_k]:
chunk = self.chunks[i]
results.append({
"doc_id": chunk["doc_id"],
"chunk_id": chunk["chunk_id"],
"text": chunk["text"],
"score": float(score),
"metadata": {
"title": chunk["title"],
"category": chunk["category"],
"source": "offline",
},
})
return results
def get_document(self, doc_id: str) -> Dict[str, Any]:
"""Return the full source document for ``doc_id``."""
doc = self.documents.get(doc_id)
if not doc:
return {"error": f"Document {doc_id} not found"}
return {
"doc_id": doc_id,
"content": doc["content"],
"metadata": {
"title": doc["title"],
"category": doc["category"],
"file": doc["file"],
"source": "offline",
},
}
quickstart.py¶
#!/usr/bin/env python3
"""Quick start script for Agentic RAG system"""
import os
import sys
import json
from pathlib import Path
def check_environment():
"""Check if environment is properly configured"""
print("🔍 Checking environment...")
# Check for .env file
if not Path(".env").exists() and Path(".env.example").exists():
print("📝 Creating .env from .env.example")
import shutil
shutil.copy(".env.example", ".env")
print("⚠️ Please edit .env and add your API keys")
return False
# Load environment variables
from dotenv import load_dotenv
load_dotenv()
# Check for at least one API key
providers = ["MOONSHOT_API_KEY", "ARK_API_KEY", "SILICONFLOW_API_KEY",
"OPENAI_API_KEY", "OPENROUTER_API_KEY"]
has_key = False
for provider in providers:
if os.getenv(provider):
has_key = True
print(f"✅ Found {provider}")
break
if not has_key:
print("❌ No API keys found. Please set at least one in .env file:")
print(" - MOONSHOT_API_KEY for Kimi")
print(" - ARK_API_KEY for Doubao")
print(" - SILICONFLOW_API_KEY for SiliconFlow")
print(" - OPENAI_API_KEY for OpenAI")
return False
return True
def setup_demo_documents():
"""Create demo documents if they don't exist"""
print("\n📚 Setting up demo documents...")
eval_dir = Path("evaluation")
eval_dir.mkdir(exist_ok=True)
# Check if documents already exist
doc_file = eval_dir / "legal_documents.json"
dataset_file = eval_dir / "legal_qa_dataset.json"
if not doc_file.exists() or not dataset_file.exists():
print("📄 Generating legal documents and dataset...")
os.chdir("evaluation")
os.system("python dataset_builder.py")
os.chdir("..")
print("✅ Documents generated")
else:
print("✅ Documents already exist")
return doc_file, dataset_file
def check_retrieval_pipeline():
"""Check if local retrieval pipeline is running"""
print("\n🔌 Checking retrieval pipeline...")
kb_type = os.getenv("KB_TYPE", "local")
if kb_type == "local":
import requests
try:
response = requests.get("http://localhost:4242/health", timeout=2)
if response.status_code == 200:
print("✅ Local retrieval pipeline is running")
return True
except:
pass
print("⚠️ Local retrieval pipeline is not running")
print(" Please run in another terminal:")
print(" cd ../retrieval-pipeline && python main.py")
print("\n Or use Dify by setting KB_TYPE=dify in .env")
return False
else:
print(f"✅ Using {kb_type} knowledge base")
return True
def index_documents(doc_file):
"""Index documents into knowledge base"""
print("\n📝 Indexing documents...")
# Check if already indexed
store_file = Path("document_store.json")
if store_file.exists():
with open(store_file, 'r', encoding='utf-8') as f:
store = json.load(f)
if len(store) > 0:
print(f"✅ Found {len(store)} documents already indexed")
return True
print("🔄 Indexing legal documents...")
result = os.system(f"python chunking.py {doc_file}")
if result == 0:
print("✅ Documents indexed successfully")
return True
else:
print("❌ Failed to index documents")
return False
def run_demo():
"""Run interactive demo"""
print("\n" + "="*60)
print("🚀 Starting Agentic RAG Demo")
print("="*60)
print("\nDemo queries you can try:")
print("1. 故意杀人罪判几年？")
print("2. 盗窃罪的立案标准是什么？")
print("3. 醉酒驾驶如何处罚？")
print("4. 张某持刀入室抢劫并造成他人重伤，应如何定罪量刑？")
print("\nCommands:")
print("- 'mode' to switch between agentic/non-agentic")
print("- 'clear' to clear conversation history")
print("- 'quit' to exit")
print("\nStarting in interactive mode...")
print("-"*60)
os.system("python main.py")
def run_comparison_demo():
"""Run comparison between agentic and non-agentic modes"""
print("\n" + "="*60)
print("🔄 Running Comparison Demo")
print("="*60)
queries = [
"故意杀人罪判几年？",
"张某因经济纠纷持刀闯入李某家中，刺伤李某致重伤并拿走5万元现金，应如何定罪？"
]
for query in queries:
print(f"\n📝 Query: {query}")
os.system(f'python main.py --mode compare --query "{query}"')
input("\nPress Enter to continue...")
def main():
"""Main quickstart function"""
print("🎯 Agentic RAG System - Quick Start")
print("="*60)
# Check environment
if not check_environment():
print("\n❌ Please configure your environment first")
sys.exit(1)
# Setup demo documents
doc_file, dataset_file = setup_demo_documents()
# Check retrieval pipeline
if not check_retrieval_pipeline():
print("\n⚠️ Warning: Retrieval pipeline not available")
print(" The system may not work properly")
response = input("\nContinue anyway? (y/n): ")
if response.lower() != 'y':
sys.exit(0)
# Index documents
if not index_documents(doc_file):
print("\n❌ Failed to index documents")
sys.exit(1)
# Menu
print("\n" + "="*60)
print("📋 Select an option:")
print("="*60)
print("1. Interactive Demo (chat with the system)")
print("2. Comparison Demo (see agentic vs non-agentic)")
print("3. Run Full Evaluation")
print("4. Exit")
choice = input("\nYour choice (1-4): ")
if choice == "1":
run_demo()
elif choice == "2":
run_comparison_demo()
elif choice == "3":
print("\n📊 Running full evaluation...")
os.chdir("evaluation")
os.system("python evaluate.py")
os.chdir("..")
elif choice == "4":
print("\n👋 Goodbye!")
else:
print("\n❌ Invalid choice")
if __name__ == "__main__":
# Install dependencies if needed
try:
import openai
import requests
from dotenv import load_dotenv
except ImportError:
print("📦 Installing required packages...")
os.system("pip install -r requirements.txt")
print("✅ Packages installed")
main()
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)")
test_structured_backends.py¶
"""
Test script for Agentic RAG with structured index backends (RAPTOR and GraphRAG).
"""
import asyncio
import logging
from config import Config, KnowledgeBaseType
from agent import AgenticRAG
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def test_raptor_backend():
"""Test Agentic RAG with RAPTOR tree-based backend"""
print("\n" + "="*60)
print("Testing RAPTOR Tree-Based Backend")
print("="*60)
# Configure for RAPTOR
config = Config.from_env()
config.knowledge_base.type = KnowledgeBaseType.RAPTOR
config.knowledge_base.raptor_base_url = "http://localhost:4242"
config.knowledge_base.raptor_top_k = 5
config.llm.provider = "kimi" # Use your preferred provider
# Initialize agent
agent = AgenticRAG(config)
# Test queries
test_queries = [
"What are the x86 general-purpose registers?",
"How does the MOV instruction work in Intel architecture?",
"Explain SIMD instructions and their purpose",
"What are control registers CR0-CR4 used for?",
"How do I use SSE instructions for parallel processing?"
]
for query in test_queries:
print(f"\nQuery: {query}")
print("-" * 50)
# Agentic mode (with tools)
response = agent.query(query, stream=False)
print(f"Response: {response[:500]}..." if len(response) > 500 else f"Response: {response}")
# Clear history for next query
agent.clear_history()
def test_graphrag_backend():
"""Test Agentic RAG with GraphRAG knowledge graph backend"""
print("\n" + "="*60)
print("Testing GraphRAG Knowledge Graph Backend")
print("="*60)
# Configure for GraphRAG
config = Config.from_env()
config.knowledge_base.type = KnowledgeBaseType.GRAPHRAG
config.knowledge_base.graphrag_base_url = "http://localhost:4242"
config.knowledge_base.graphrag_top_k = 5
config.knowledge_base.graphrag_search_type = "hybrid"
config.llm.provider = "kimi" # Use your preferred provider
# Initialize agent
agent = AgenticRAG(config)
# Test queries
test_queries = [
"What instructions modify the FLAGS register?",
"Show me the relationship between MOV and LEA instructions",
"What CPU features are related to virtualization?",
"How are SSE and AVX instructions related?",
"What components make up the execution environment?"
]
for query in test_queries:
print(f"\nQuery: {query}")
print("-" * 50)
# Agentic mode (with tools)
response = agent.query(query, stream=False)
print(f"Response: {response[:500]}..." if len(response) > 500 else f"Response: {response}")
# Clear history for next query
agent.clear_history()
def compare_backends():
"""Compare results from different backends for the same query"""
print("\n" + "="*60)
print("Comparing Different Backend Results")
print("="*60)
query = "Explain the Intel x86 instruction format and its components"
backends = [
(KnowledgeBaseType.RAPTOR, "RAPTOR Tree-Based", "http://localhost:4242"),
(KnowledgeBaseType.GRAPHRAG, "GraphRAG Knowledge Graph", "http://localhost:4242")
]
results = {}
for backend_type, backend_name, base_url in backends:
print(f"\n{backend_name}:")
print("-" * 40)
# Configure for backend
config = Config.from_env()
config.knowledge_base.type = backend_type
if backend_type == KnowledgeBaseType.RAPTOR:
config.knowledge_base.raptor_base_url = base_url
elif backend_type == KnowledgeBaseType.GRAPHRAG:
config.knowledge_base.graphrag_base_url = base_url
config.knowledge_base.graphrag_search_type = "hybrid"
config.llm.provider = "kimi"
# Initialize agent
agent = AgenticRAG(config)
# Query and store result
response = agent.query(query, stream=False)
results[backend_name] = response
print(f"Response preview: {response[:300]}...")
# Compare results
print("\n" + "="*60)
print("Comparison Summary")
print("="*60)
for backend_name, response in results.items():
print(f"\n{backend_name}:")
print(f" Response length: {len(response)} characters")
print(f" Citations found: {'[Doc:' in response or '[Chunk:' in response}")
# Count tool calls (approximate)
tool_indicators = ["knowledge_base_search", "get_document"]
tool_count = sum(1 for indicator in tool_indicators if indicator in str(response))
print(f" Estimated tool calls: {tool_count}")
def test_non_agentic_mode():
"""Test non-agentic mode with structured backends"""
print("\n" + "="*60)
print("Testing Non-Agentic Mode with Structured Backends")
print("="*60)
query = "What are the different types of Intel CPU registers?"
# Test with RAPTOR
print("\nRAPTOR (Non-Agentic):")
config = Config.from_env()
config.knowledge_base.type = KnowledgeBaseType.RAPTOR
config.knowledge_base.raptor_base_url = "http://localhost:4242"
agent = AgenticRAG(config)
response = agent.query_non_agentic(query, stream=False)
print(f"Response: {response[:400]}...")
# Test with GraphRAG
print("\nGraphRAG (Non-Agentic):")
config.knowledge_base.type = KnowledgeBaseType.GRAPHRAG
config.knowledge_base.graphrag_base_url = "http://localhost:4242"
agent = AgenticRAG(config)
response = agent.query_non_agentic(query, stream=False)
print(f"Response: {response[:400]}...")
def main():
"""Run all tests"""
print("Agentic RAG with Structured Index Backends Test Suite")
print("=" * 60)
# Make sure the structured-index API is running on port 4242
print("\nNote: Make sure the structured-index API is running on port 4242")
print("Run: cd ../structured-index && python main.py serve")
input("\nPress Enter to start tests...")
try:
# Run tests
test_raptor_backend()
test_graphrag_backend()
compare_backends()
test_non_agentic_mode()
print("\n" + "="*60)
print("All tests completed successfully!")
print("="*60)
except Exception as e:
logger.error(f"Test failed: {e}")
print("\nMake sure:")
print("1. The structured-index API is running (python main.py serve)")
print("2. Indexes have been built (python main.py build <document>)")
print("3. Your API keys are configured in .env")
if __name__ == "__main__":
main()
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
self._offline_retriever = None # Lazily built in-process BM25 index
# 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:
if self.config.type == KnowledgeBaseType.OFFLINE:
return self._search_offline(query)
elif 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:
raise ValueError(f"Unsupported knowledge base type: {self.config.type}")
except Exception as e:
logger.error(f"Error in knowledge base search: {e}")
return []
def _get_offline_retriever(self):
"""Lazily build the in-process BM25 retriever over the local corpus."""
if self._offline_retriever is None:
from offline_retriever import OfflineRetriever
self._offline_retriever = OfflineRetriever(self.config.offline_corpus_path)
return self._offline_retriever
def _search_offline(self, query: str) -> List[Dict[str, Any]]:
"""Search the in-process BM25 index (no server, no API key required)."""
retriever = self._get_offline_retriever()
results = retriever.search(query, self.config.offline_top_k)
logger.info(f"Offline BM25 search returned {len(results)} results")
return results
def _search_local(self, query: str) -> List[Dict[str, Any]]:
"""Search using local retrieval pipeline"""
try:
response = requests.post(
f"{self.config.local_base_url}/search",
json={
"query": query,
"mode": "hybrid",
"top_k": self.config.local_top_k,
"rerank": True
}
)
response.raise_for_status()
results = []
data = response.json()
# The retrieval pipeline returns results in different keys based on mode
# For hybrid mode, we want the reranked_results
search_results = data.get("reranked_results", [])
# If no reranked results, fall back to dense or sparse results
if not search_results:
search_results = data.get("dense_results", [])
if not search_results:
search_results = data.get("sparse_results", [])
for item in search_results:
# Extract doc_id and chunk_id from the result
doc_id = item.get("doc_id", "")
chunk_id = item.get("chunk_id", f"{doc_id}_chunk_{len(results)}")
# Get the text field and score based on result type
text = item.get("text", "")
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")
return results
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.OFFLINE:
return self._get_offline_retriever().get_document(doc_id)
elif 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/offline_qa.json¶
{
"description": "离线检索对比数据集：用于在无需 LLM / 无需外部检索服务的情况下，量化对比『非智能体化 RAG（单次检索）』与『智能体化 RAG（多轮/分解检索）』的证据召回能力。每条问题标注了回答所必需的金标准法条（gold_articles，以法条编号精确匹配）。naive_query 为用户原始提问（单次检索直接使用）；subqueries 为智能体经过思考后分解/改写出的检索式（多次检索后取并集）。金标准法条均已确认存在于 laws/ 语料中。",
"metric": "证据召回率 = 命中的金标准法条数 / 金标准法条总数（某法条被命中当且仅当检索结果中存在以该法条编号开头的分块）",
"cases": [
{
"id": "easy_1",
"question": "故意伤害致人重伤的，如何处罚？",
"difficulty": "easy",
"naive_query": "故意伤害致人重伤的，如何处罚？",
"subqueries": ["故意伤害 致人重伤 有期徒刑"],
"gold_articles": ["第二百三十四条"]
},
{
"id": "easy_2",
"question": "正当防卫是怎么规定的？",
"difficulty": "easy",
"naive_query": "正当防卫是怎么规定的？",
"subqueries": ["正当防卫 不负刑事责任"],
"gold_articles": ["第二十条"]
},
{
"id": "easy_3",
"question": "醉酒驾驶机动车如何处罚？",
"difficulty": "easy",
"naive_query": "醉酒驾驶机动车如何处罚？",
"subqueries": ["危险驾驶罪 醉酒 驾驶机动车 拘役"],
"gold_articles": ["第一百三十三条之一"]
},
{
"id": "hard_1",
"question": "故意杀人罪判几年？",
"difficulty": "hard",
"naive_query": "故意杀人罪判几年？",
"subqueries": ["故意杀人 死刑 无期徒刑"],
"gold_articles": ["第二百三十二条"]
},
{
"id": "hard_2",
"question": "盗窃罪的立案标准是什么？",
"difficulty": "hard",
"naive_query": "盗窃罪的立案标准是什么？",
"subqueries": ["盗窃 公私财物 数额较大 有期徒刑"],
"gold_articles": ["第二百六十四条"]
},
{
"id": "hard_3",
"question": "诈骗罪的量刑标准是什么？",
"difficulty": "hard",
"naive_query": "诈骗罪的量刑标准是什么？",
"subqueries": ["诈骗 公私财物 数额较大 有期徒刑 罚金"],
"gold_articles": ["第二百六十六条"]
},
{
"id": "hard_4",
"question": "醉酒过失致人重伤且有盗窃前科，应如何量刑？",
"difficulty": "hard",
"naive_query": "醉酒过失致人重伤且有盗窃前科，应如何量刑？",
"subqueries": [
"过失伤害他人致人重伤 有期徒刑",
"危险驾驶罪 醉酒 驾驶机动车",
"累犯 从重处罚 过失犯罪"
],
"gold_articles": ["第二百三十五条", "第一百三十三条之一", "第六十五条"]
}
]
}