ai-agent-book 精选快照(<2MB 代码与文档,来自 github.com/bojieli/ai-agent-book)
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
Build latest book artifacts / build (push) Canceled after 0s
dependency resolution / resolve (3.11) (push) Canceled after 0s
dependency resolution / resolve (3.13) (push) Canceled after 0s
deploy-pages / build (push) Canceled after 0s
deploy-pages / deploy (push) Canceled after 0s
i18n consistency check / check (push) Canceled after 0s
provider adoption tests / test (chapter2/context-compression) (push) Canceled after 0s
provider adoption tests / test (chapter2/prompt-injection) (push) Canceled after 0s
provider adoption tests / test (chapter2/system-hint) (push) Canceled after 0s
provider adoption tests / test (chapter3/log-sanitization) (push) Canceled after 0s
web-search-agent tests / test (push) Canceled after 0s
web-search-agent tests / agentbook (push) Canceled after 0s
This commit is contained in:
@@ -0,0 +1,251 @@
|
||||
# Contextual Retrieval System / 上下文感知检索系统
|
||||
|
||||
> Companion material for *AI Agents in Depth*, Chapter 3 — **Experiment 3-10**: Anthropic-style contextual prefixes before indexing; offline BM25 recall@k compare.
|
||||
> 配套《深入理解 AI Agent》第 3 章 **实验 3-10**:索引前为分块生成上下文前缀;离线 BM25 recall@k 对比。
|
||||
|
||||
← [Chapter 3 index / 返回第 3 章目录](../README.md)
|
||||
|
||||
---
|
||||
|
||||
## English
|
||||
|
||||
### Canonical live campaign
|
||||
|
||||
`python campaign.py` regenerates every prefix live from the full source
|
||||
document plus its target chunk, then compares the same 15 labeled queries over
|
||||
plain/contextual BM25, dense Qwen3 embeddings, and RRF hybrid indexes. It records
|
||||
recall@1/3/5, MRR, model revision, index-time usage/cost assumptions, source
|
||||
hashes, and raw ARK receipts in `validation/runs/<run-id>/`; the canonical
|
||||
pointer is `validation/latest.json`. `compare_retrieval.py` remains the small
|
||||
offline demonstration and is not accepted as the canonical live result.
|
||||
|
||||
### Overview
|
||||
|
||||
Educational implementation of Anthropic’s Contextual Retrieval: prepend chunk-specific context before embedding/indexing to fix the “orphaned chunk” problem.
|
||||
|
||||
### Key insight
|
||||
|
||||
**Problem:** Traditional RAG loses context when chunking. “The company’s revenue grew by 3%” is meaningless without which company / which period.
|
||||
|
||||
**Solution:** Generate short explanatory context per chunk and prepend it before indexing so BM25 and embeddings keep identity signals.
|
||||
|
||||
### Core offline experiment (Experiment 3-10)
|
||||
|
||||
`compare_retrieval.py` quantifies the claim **fully offline**: same chunks indexed two ways—plain (raw text only) vs contextual (LLM-generated prefix + text)—then compares `recall@k` on `evaluation/retrieval_eval.json` (15 queries + human gold chunks). **No API or retrieval service** (BM25 + jieba).
|
||||
|
||||
```bash
|
||||
python compare_retrieval.py
|
||||
python compare_retrieval.py --per-query
|
||||
python compare_retrieval.py --query "国家主席有哪些职权?" --top-k 5
|
||||
python compare_retrieval.py --mode plain
|
||||
python compare_retrieval.py --output result.json
|
||||
python compare_retrieval.py --help # Chinese help
|
||||
```
|
||||
|
||||
Real run (22 Constitution / Prosecutor Law chunks, 15 queries, jieba):
|
||||
|
||||
```
|
||||
检索召回对比:无上下文分块 vs. 上下文感知检索(BM25)
|
||||
====================================================================
|
||||
方法 recall@1 recall@3 recall@5
|
||||
----------------------------------------------------
|
||||
无上下文 (plain) 60.0% 86.7% 93.3%
|
||||
有上下文 (ctx) 86.7% 86.7% 93.3%
|
||||
----------------------------------------------------
|
||||
提升 (Δpp) +26.7pp +0.0pp +0.0pp
|
||||
----------------------------------------------------
|
||||
失败率下降 67% 0% 0%
|
||||
```
|
||||
|
||||
Conclusion (matches the book): context prefixes lift top-1 recall (60% → 86.7%; failure rate 1−recall@1 down 67%). Gain is strongest at recall@1; `--query` shows how the prefix re-ranks the correct section first.
|
||||
|
||||
> `--method embedding` / `--method hybrid` need embedding APIs (not offline); the script falls back to BM25 offline results. Full dense + rerank lives in `contextual_tools.py`.
|
||||
> Same logic is also in `ContextualChunker.compare_retrieval_methods()`.
|
||||
|
||||
### Educational features
|
||||
|
||||
1. Watch LLM context generation per chunk
|
||||
2. Dual indexing (BM25 + embeddings) benefits from context
|
||||
3. Compare with `use_contextual=False`
|
||||
4. Metrics and token/cost awareness
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
Document → Basic Chunking → Context Generation (optional LLM)
|
||||
→ Enhanced chunks (context+text vs text only)
|
||||
→ Retrieval pipeline (sparse BM25 + dense embeddings)
|
||||
→ Hybrid search + reranking
|
||||
```
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
# From the repository root: use the shared Chapter 3 environment
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# Activate it before changing directories:
|
||||
# macOS/Linux:
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell: .venv\Scripts\Activate.ps1
|
||||
# Windows cmd: .venv\Scripts\activate.bat
|
||||
|
||||
# pip fallback when uv is not installed:
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/contextual-retrieval
|
||||
|
||||
# Single-project compatibility path, still supported during migration:
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
# MOONSHOT_API_KEY / ARK_API_KEY / OPENAI_API_KEY / etc.
|
||||
|
||||
# Separate terminal for full e2e with pipeline:
|
||||
cd ../retrieval-pipeline
|
||||
python main.py
|
||||
# http://localhost:4242
|
||||
|
||||
# Back in this project directory (or a second terminal):
|
||||
cd ../contextual-retrieval
|
||||
|
||||
# Index with contextual enhancement
|
||||
python index_local_laws_contextual.py
|
||||
python index_local_laws_contextual.py --no-contextual
|
||||
|
||||
# Queries
|
||||
python main.py
|
||||
python main.py --query "宪法第一条是什么" --mode agentic
|
||||
python main.py --query "宪法第一条是什么" --mode compare
|
||||
```
|
||||
|
||||
### Context generation process
|
||||
|
||||
1. Provide full document (or surrounding context) to the LLM
|
||||
2. Show the specific chunk
|
||||
3. Ask for 2–3 sentence situating context
|
||||
|
||||
Template sketch:
|
||||
|
||||
```
|
||||
<document>
|
||||
[Full document or surrounding context]
|
||||
</document>
|
||||
|
||||
Here is the chunk we want to situate:
|
||||
<chunk>
|
||||
[Specific chunk text]
|
||||
</chunk>
|
||||
|
||||
Please give a short, succinct context to situate this chunk within the overall document...
|
||||
```
|
||||
|
||||
### References / license
|
||||
|
||||
- [Anthropic Contextual Retrieval](https://www.anthropic.com/engineering/contextual-retrieval)
|
||||
- Educational project for learning purposes. Acknowledgments: Anthropic engineering research.
|
||||
|
||||
---
|
||||
|
||||
## 中文
|
||||
|
||||
### 概述
|
||||
|
||||
Anthropic 上下文感知检索的教学实现:在嵌入/建索引前为每个分块附加专属上下文,缓解「孤儿分块」问题。
|
||||
|
||||
### 核心洞察
|
||||
|
||||
**问题:** 传统 RAG 分块后丢失语境。「公司收入增长 3%」不知是哪家公司、哪个时期。
|
||||
**方案:** 为每块生成简短解释性上下文并前置,使 BM25 与向量都能保留身份信号。
|
||||
|
||||
### 核心实验:离线量化召回提升(实验 3-10)
|
||||
|
||||
`compare_retrieval.py` **完全离线**量化:同一批文本块分别以无上下文 / 有上下文前缀两种方式建 BM25 索引,在 `evaluation/retrieval_eval.json`(15 查询 + 人工 gold 块)上比较 `recall@k`。**无需 API 或检索服务**(BM25 + jieba)。
|
||||
|
||||
```bash
|
||||
python compare_retrieval.py
|
||||
python compare_retrieval.py --per-query
|
||||
python compare_retrieval.py --query "国家主席有哪些职权?" --top-k 5
|
||||
python compare_retrieval.py --mode plain
|
||||
python compare_retrieval.py --output result.json
|
||||
python compare_retrieval.py --help
|
||||
```
|
||||
|
||||
真实输出(22 个《宪法》《检察官法》文本块,15 查询):
|
||||
|
||||
```
|
||||
检索召回对比:无上下文分块 vs. 上下文感知检索(BM25)
|
||||
====================================================================
|
||||
方法 recall@1 recall@3 recall@5
|
||||
----------------------------------------------------
|
||||
无上下文 (plain) 60.0% 86.7% 93.3%
|
||||
有上下文 (ctx) 86.7% 86.7% 93.3%
|
||||
----------------------------------------------------
|
||||
提升 (Δpp) +26.7pp +0.0pp +0.0pp
|
||||
----------------------------------------------------
|
||||
失败率下降 67% 0% 0%
|
||||
```
|
||||
|
||||
结论与书中一致:上下文前缀显著提升 top-1 召回(60% → 86.7%,失败率下降 67%)。`--method embedding` / `hybrid` 需 embedding API,脚本会提示并回退 BM25 离线结果。完整稠密+重排见 `contextual_tools.py`。
|
||||
|
||||
### 教学特性
|
||||
|
||||
观察上下文生成、双索引策略、`use_contextual=False` 对照、指标与成本。
|
||||
|
||||
### 架构
|
||||
|
||||
文档 → 基础分块 → 可选 LLM 上下文生成 → 增强块 → 检索流水线(稀疏+稠密)→ 混合检索+重排。
|
||||
|
||||
### 快速开始
|
||||
|
||||
```bash
|
||||
# 在仓库根目录使用统一的第 3 章环境
|
||||
uv sync --locked --python 3.12 --extra ch3
|
||||
|
||||
# 切换目录前先激活环境:
|
||||
# macOS/Linux:
|
||||
source .venv/bin/activate
|
||||
# Windows PowerShell:.venv\Scripts\Activate.ps1
|
||||
# Windows cmd:.venv\Scripts\activate.bat
|
||||
|
||||
# 未安装 uv 时可用 pip 兜底:
|
||||
# python -m pip install -e ".[ch3]"
|
||||
|
||||
cd chapter3/contextual-retrieval
|
||||
|
||||
# 迁移期间仍支持单项目兼容路径:
|
||||
# python -m pip install -r requirements.txt
|
||||
|
||||
cp env.example .env
|
||||
|
||||
# 另开一个终端运行完整检索流水线:
|
||||
cd ../retrieval-pipeline
|
||||
python main.py
|
||||
# http://localhost:4242
|
||||
|
||||
# 回到本项目目录(或新开一个终端):
|
||||
cd ../contextual-retrieval
|
||||
|
||||
python index_local_laws_contextual.py
|
||||
python index_local_laws_contextual.py --no-contextual
|
||||
|
||||
python main.py
|
||||
python main.py --query "宪法第一条是什么" --mode agentic
|
||||
python main.py --query "宪法第一条是什么" --mode compare
|
||||
```
|
||||
|
||||
### 上下文生成流程
|
||||
|
||||
向 LLM 提供全文(或周边)+ 目标块,请求 2–3 句定位上下文。提示模板见 English 节。
|
||||
|
||||
### 参考与许可
|
||||
|
||||
[Anthropic 博客](https://www.anthropic.com/engineering/contextual-retrieval) · 教学项目。
|
||||
|
||||
---
|
||||
|
||||
## Notes / 说明
|
||||
|
||||
### OpenRouter 通用回退 / Universal OpenRouter fallback
|
||||
|
||||
Chat LLM falls back to OpenRouter when primary keys are missing and `OPENROUTER_API_KEY` is set. See `env.example`. Related: [`../contextual-retrieval-for-user-memory/`](../contextual-retrieval-for-user-memory/) (Exp. 3-11).
|
||||
@@ -0,0 +1,89 @@
|
||||
# Contextual Legal Document Indexing
|
||||
|
||||
This script implements Anthropic's Contextual Retrieval approach for indexing Chinese legal documents.
|
||||
|
||||
## Key Innovation: Contextual Retrieval
|
||||
|
||||
Unlike traditional RAG that loses context when chunking, this script:
|
||||
1. Generates contextual descriptions for each chunk using LLM
|
||||
2. Prepends context to chunks before indexing
|
||||
3. Significantly improves retrieval accuracy
|
||||
|
||||
## Features
|
||||
|
||||
- **Contextual Enhancement**: Uses LLM to generate chunk-specific context
|
||||
- **Smart Chunking**: Paragraph-aware boundaries (soft: 1024, hard: 2048 chars)
|
||||
- **Comparison Mode**: Run with/without context for performance comparison
|
||||
- **Cache Optimization**: Caches context for similar chunks to reduce API costs
|
||||
- **Detailed Statistics**: Token usage, generation time, and cost estimation
|
||||
|
||||
## Prerequisites
|
||||
|
||||
1. Set up your LLM API key:
|
||||
```bash
|
||||
export MOONSHOT_API_KEY="your_api_key" # Default: Kimi
|
||||
# Or use other providers:
|
||||
export OPENAI_API_KEY="your_api_key"
|
||||
export SILICONFLOW_API_KEY="your_api_key"
|
||||
```
|
||||
|
||||
2. Ensure retrieval pipeline is running:
|
||||
```bash
|
||||
# Terminal 1: Dense service
|
||||
python dense_service.py
|
||||
|
||||
# Terminal 2: Sparse service
|
||||
python sparse_service.py
|
||||
|
||||
# Terminal 3: Main pipeline
|
||||
python main.py
|
||||
```
|
||||
|
||||
3. The `laws` directory should be linked/present (automatically created as symlink to agentic-rag/laws)
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Contextual Indexing
|
||||
```bash
|
||||
# Index with contextual enhancement (default)
|
||||
python index_local_laws_contextual.py
|
||||
```
|
||||
|
||||
### Advanced Options
|
||||
```bash
|
||||
# Process limited documents
|
||||
python index_local_laws_contextual.py --max-docs 10
|
||||
|
||||
# Process specific categories
|
||||
python index_local_laws_contextual.py --categories "宪法" "民法典"
|
||||
|
||||
# Use different LLM provider
|
||||
python index_local_laws_contextual.py --llm-provider openai --llm-model gpt-5.6-luna
|
||||
|
||||
# Custom batch size for indexing
|
||||
python index_local_laws_contextual.py --batch-size 20
|
||||
|
||||
# Skip cleanup
|
||||
python index_local_laws_contextual.py --no-cleanup
|
||||
```
|
||||
|
||||
## Cost Considerations
|
||||
|
||||
Context generation requires LLM API calls:
|
||||
- ~150 tokens per chunk for context generation
|
||||
- Costs vary by provider (OpenAI: ~$0.03/1K tokens, Others: ~$0.01/1K tokens)
|
||||
- Cache reduces costs for duplicate content
|
||||
|
||||
Estimate for 288 legal documents:
|
||||
- ~3000-5000 chunks total
|
||||
- ~450K-750K tokens
|
||||
- Cost: $5-15 depending on provider
|
||||
|
||||
## Document Store
|
||||
|
||||
Maintains `document_store.json` with:
|
||||
- Document metadata
|
||||
- Chunk statistics
|
||||
- Context token usage
|
||||
- Generation metrics
|
||||
- Indexing timestamps
|
||||
@@ -0,0 +1,100 @@
|
||||
# Semantic Document and Chunk IDs
|
||||
|
||||
## Overview
|
||||
|
||||
The contextual retrieval system now uses semantically meaningful document IDs based on file names instead of opaque MD5 hashes. This makes the system more transparent, debuggable, and user-friendly.
|
||||
|
||||
## ID Generation Rules
|
||||
|
||||
### Document ID
|
||||
Generated from the file name with these transformations:
|
||||
1. Remove file extension (.md)
|
||||
2. Replace Chinese parentheses () with underscores
|
||||
3. Replace English parentheses () with underscores
|
||||
4. Replace spaces and hyphens with underscores
|
||||
5. Remove trailing underscores
|
||||
6. Truncate to 100 characters if needed
|
||||
|
||||
### Chunk ID
|
||||
Format: `{document_id}_chunk_{index}`
|
||||
- Document ID as base
|
||||
- Sequential chunk index (0, 1, 2...)
|
||||
|
||||
## Examples
|
||||
|
||||
| Original File | Document ID | Sample Chunk IDs |
|
||||
|--------------|-------------|------------------|
|
||||
| 宪法.md | 宪法 | 宪法_chunk_0, 宪法_chunk_1 |
|
||||
| 劳动法(2018-12-29).md | 劳动法_2018_12_29 | 劳动法_2018_12_29_chunk_0 |
|
||||
| 民法典总则编.md | 民法典总则编 | 民法典总则编_chunk_0 |
|
||||
| 检察官法(2019-04-23).md | 检察官法_2019_04_23 | 检察官法_2019_04_23_chunk_0 |
|
||||
|
||||
## Comparison with Hash-Based IDs
|
||||
|
||||
### Old System (MD5 Hash)
|
||||
```
|
||||
Document: 08f758bf19c0
|
||||
Chunks: 08f758bf19c0_chunk_0, 08f758bf19c0_chunk_1
|
||||
```
|
||||
- ❌ Not human-readable
|
||||
- ❌ No semantic meaning
|
||||
- ❌ Hard to debug
|
||||
- ❌ Can't identify source document
|
||||
|
||||
### New System (Semantic)
|
||||
```
|
||||
Document: 宪法
|
||||
Chunks: 宪法_chunk_0, 宪法_chunk_1
|
||||
```
|
||||
- ✅ Human-readable
|
||||
- ✅ Self-documenting
|
||||
- ✅ Easy to debug
|
||||
- ✅ Clear source identification
|
||||
|
||||
## Benefits
|
||||
|
||||
1. **Transparency**: Users and developers can immediately identify which document a chunk comes from
|
||||
2. **Searchability**: Can grep/search for specific laws by name in logs and data
|
||||
3. **Debugging**: Easier to trace issues back to source documents
|
||||
4. **Consistency**: Same document always generates the same ID
|
||||
5. **Sortability**: Documents sort alphabetically by name
|
||||
|
||||
## Implementation
|
||||
|
||||
The ID generation is handled by the `generate_document_id()` method in `index_local_laws_contextual.py`:
|
||||
|
||||
```python
|
||||
def generate_document_id(self, doc_info: Dict[str, Any]) -> str:
|
||||
"""Generate a semantically meaningful document ID from file name."""
|
||||
base_name = doc_info["name"]
|
||||
|
||||
# Clean up the name
|
||||
clean_name = base_name.replace('(', '_').replace(')', '')
|
||||
clean_name = clean_name.replace('(', '_').replace(')', '')
|
||||
clean_name = re.sub(r'[\s\-]+', '_', clean_name)
|
||||
clean_name = clean_name.strip('_')
|
||||
|
||||
return clean_name
|
||||
```
|
||||
|
||||
## Testing
|
||||
|
||||
Run the test script to see examples:
|
||||
```bash
|
||||
python test_document_ids.py
|
||||
```
|
||||
|
||||
## Migration Note
|
||||
|
||||
If you have existing indexed documents with hash-based IDs, you'll need to re-index them to use the new semantic IDs:
|
||||
|
||||
```bash
|
||||
python index_local_laws_contextual.py
|
||||
```
|
||||
|
||||
## Future Enhancements
|
||||
|
||||
Potential improvements for the ID generation:
|
||||
1. Add version tracking (e.g., 劳动法_v2018_12_29)
|
||||
2. Include document type prefix (e.g., law_劳动法, regulation_xxx)
|
||||
3. Support for hierarchical documents (e.g., 民法典/总则编 → 民法典_总则编)
|
||||
@@ -0,0 +1,390 @@
|
||||
"""Agentic RAG System with ReAct Pattern"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional, Generator
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from openai import OpenAI
|
||||
|
||||
from config import Config, LLMConfig, AgentConfig
|
||||
from tools import KnowledgeBaseTools, get_tool_definitions
|
||||
|
||||
|
||||
def _reasoning_safe_temperature(model, requested=1.0):
|
||||
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
|
||||
Return 1 for those; otherwise the requested value so non-reasoning
|
||||
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
|
||||
m = str(model or "").lower().replace("/", "-")
|
||||
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
|
||||
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Message:
|
||||
"""Represents a message in the conversation"""
|
||||
role: str # "user", "assistant", "tool"
|
||||
content: str
|
||||
tool_calls: Optional[List[Dict[str, Any]]] = None
|
||||
tool_call_id: Optional[str] = None
|
||||
timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
|
||||
|
||||
|
||||
class AgenticRAG:
|
||||
"""Agentic RAG system with ReAct pattern and multiple LLM provider support"""
|
||||
|
||||
def __init__(self, config: Optional[Config] = None):
|
||||
"""Initialize the agent"""
|
||||
self.config = config or Config.from_env()
|
||||
|
||||
# Initialize LLM client
|
||||
self._init_llm_client()
|
||||
|
||||
# Initialize knowledge base tools
|
||||
self.kb_tools = KnowledgeBaseTools(self.config.knowledge_base)
|
||||
|
||||
# Conversation history
|
||||
self.conversation_history: List[Dict[str, Any]] = []
|
||||
|
||||
# Tool definitions
|
||||
self.tools = get_tool_definitions()
|
||||
|
||||
logger.info(f"Initialized AgenticRAG with provider: {self.config.llm.provider}")
|
||||
|
||||
def _init_llm_client(self):
|
||||
"""Initialize the LLM client based on provider"""
|
||||
client_config, model = self.config.llm.get_client_config()
|
||||
|
||||
# Extract base_url if present
|
||||
base_url = client_config.pop("base_url", None)
|
||||
|
||||
# Create OpenAI client
|
||||
if base_url:
|
||||
self.client = OpenAI(base_url=base_url, **client_config)
|
||||
else:
|
||||
self.client = OpenAI(**client_config)
|
||||
|
||||
self.model = model
|
||||
logger.info(f"Using model: {self.model}")
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
"""Generate the system prompt"""
|
||||
return """You are an intelligent assistant with access to a knowledge base. Your primary role is to answer questions accurately based on the information available in the knowledge base.
|
||||
|
||||
## Important Guidelines:
|
||||
|
||||
1. **Knowledge Base Only**: You MUST only answer questions based on information found in the knowledge base. If the information is not available, clearly state that you cannot answer based on the available knowledge.
|
||||
|
||||
2. **Use Tools Effectively**:
|
||||
- Use `knowledge_base_search` to search for relevant information
|
||||
- Use `get_document` to retrieve complete documents when you need more context
|
||||
- You may need multiple searches with different queries to fully answer a question
|
||||
|
||||
3. **Citations Required**: Always include citations in your answers. Format citations as [Doc: document_id] or [Chunk: chunk_id] inline with your response.
|
||||
|
||||
4. **Reasoning Process**: Think step-by-step:
|
||||
- First, understand what information is needed
|
||||
- Search for relevant information
|
||||
- If needed, retrieve full documents for context
|
||||
- Synthesize the information to answer the question
|
||||
- Include proper citations
|
||||
|
||||
5. **Handle Follow-ups**: For follow-up questions, consider the conversation context but always verify information from the knowledge base.
|
||||
|
||||
6. **Be Accurate**: Never make up information. If something is unclear or not found, say so explicitly.
|
||||
|
||||
Remember: Your credibility depends on providing accurate, well-cited information from the knowledge base only."""
|
||||
|
||||
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""Execute a tool and return the result"""
|
||||
try:
|
||||
if tool_name == "knowledge_base_search":
|
||||
query = arguments.get("query", "")
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Executing tool: {tool_name} with args: {arguments}")
|
||||
|
||||
results = self.kb_tools.knowledge_base_search(query)
|
||||
|
||||
if not results:
|
||||
logger.info(f"No results found for query: {query}")
|
||||
return {"status": "no_results", "message": f"No relevant documents found for query: {query}"}
|
||||
|
||||
# Format results for agent
|
||||
formatted_results = []
|
||||
for r in results:
|
||||
formatted_results.append({
|
||||
"doc_id": r["doc_id"],
|
||||
"chunk_id": r["chunk_id"],
|
||||
"text": r["text"],
|
||||
"score": r["score"]
|
||||
})
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"results": formatted_results,
|
||||
"total_found": len(results)
|
||||
}
|
||||
|
||||
elif tool_name == "get_document":
|
||||
doc_id = arguments.get("doc_id", "")
|
||||
document = self.kb_tools.get_document(doc_id)
|
||||
|
||||
if "error" in document:
|
||||
return {"status": "error", "message": document["error"]}
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"document": {
|
||||
"doc_id": document.get("doc_id", doc_id),
|
||||
"content": document.get("content", ""),
|
||||
"metadata": document.get("metadata", {})
|
||||
}
|
||||
}
|
||||
|
||||
else:
|
||||
return {"status": "error", "message": f"Unknown tool: {tool_name}"}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Tool execution error: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
def _build_messages(self, user_query: str) -> List[Dict[str, Any]]:
|
||||
"""Build messages for the LLM including conversation history"""
|
||||
messages = [{"role": "system", "content": self._get_system_prompt()}]
|
||||
|
||||
# Add conversation history (limited)
|
||||
history_limit = self.config.agent.conversation_history_limit
|
||||
# limit<=0 → no history; list[-0:] would include all turns.
|
||||
if history_limit > 0:
|
||||
if len(self.conversation_history) > history_limit:
|
||||
messages.extend(self.conversation_history[-history_limit:])
|
||||
else:
|
||||
messages.extend(self.conversation_history)
|
||||
|
||||
# Add current user query
|
||||
messages.append({"role": "user", "content": user_query})
|
||||
|
||||
return messages
|
||||
|
||||
def query(self, user_query: str, stream: bool = None) -> Any:
|
||||
"""
|
||||
Process a user query using the ReAct pattern.
|
||||
|
||||
Args:
|
||||
user_query: The user's question
|
||||
stream: Whether to stream the response
|
||||
|
||||
Returns:
|
||||
The agent's response (string or generator for streaming)
|
||||
"""
|
||||
if stream is None:
|
||||
stream = self.config.llm.stream
|
||||
|
||||
# Build messages
|
||||
messages = self._build_messages(user_query)
|
||||
|
||||
# Track iterations
|
||||
iterations = 0
|
||||
max_iterations = self.config.agent.max_iterations
|
||||
|
||||
# Process with ReAct loop
|
||||
while iterations < max_iterations:
|
||||
iterations += 1
|
||||
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Iteration {iterations}/{max_iterations}")
|
||||
|
||||
try:
|
||||
# Call LLM with tools
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
tools=self.tools,
|
||||
tool_choice="auto",
|
||||
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
|
||||
max_tokens=self.config.llm.max_tokens,
|
||||
stream=False # We handle streaming separately
|
||||
)
|
||||
|
||||
message = response.choices[0].message
|
||||
|
||||
# Add assistant message to history
|
||||
assistant_msg = {"role": "assistant", "content": message.content or ""}
|
||||
if message.tool_calls:
|
||||
assistant_msg["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": tc.type,
|
||||
"function": {
|
||||
"name": tc.function.name,
|
||||
"arguments": tc.function.arguments
|
||||
}
|
||||
} for tc in message.tool_calls
|
||||
]
|
||||
messages.append(assistant_msg)
|
||||
|
||||
# Process tool calls if present
|
||||
if message.tool_calls:
|
||||
for tool_call in message.tool_calls:
|
||||
tool_name = tool_call.function.name
|
||||
try:
|
||||
arguments = json.loads(tool_call.function.arguments)
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to parse tool arguments: {tool_call.function.arguments}")
|
||||
arguments = {}
|
||||
|
||||
# Execute tool
|
||||
result = self._execute_tool(tool_name, arguments)
|
||||
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Tool result: {json.dumps(result, indent=2, ensure_ascii=False)}")
|
||||
|
||||
# Add tool result to messages
|
||||
tool_message = {
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call.id,
|
||||
"content": json.dumps(result, indent=2, ensure_ascii=False)
|
||||
}
|
||||
messages.append(tool_message)
|
||||
|
||||
# Continue loop for next iteration
|
||||
continue
|
||||
else:
|
||||
# No tool calls, we have final answer
|
||||
# Update conversation history
|
||||
self.conversation_history.append({"role": "user", "content": user_query})
|
||||
self.conversation_history.append(assistant_msg)
|
||||
|
||||
final_response = message.content or ""
|
||||
|
||||
# Log the final response if verbose
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Final response generated (length: {len(final_response)} chars)")
|
||||
|
||||
# Return response
|
||||
if stream:
|
||||
return self._stream_response(final_response)
|
||||
else:
|
||||
return final_response
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in query processing: {e}")
|
||||
error_msg = f"Error processing query: {str(e)}"
|
||||
if stream:
|
||||
return self._stream_response(error_msg)
|
||||
else:
|
||||
return error_msg
|
||||
|
||||
# Max iterations reached
|
||||
logger.warning(f"Max iterations ({max_iterations}) reached")
|
||||
final_msg = "I need more iterations to fully answer your question. Please try rephrasing or breaking down your query."
|
||||
|
||||
if stream:
|
||||
return self._stream_response(final_msg)
|
||||
else:
|
||||
return final_msg
|
||||
|
||||
def _stream_response(self, content: str) -> Generator[str, None, None]:
|
||||
"""Stream response content"""
|
||||
# Simple character streaming for demonstration
|
||||
for char in content:
|
||||
yield char
|
||||
|
||||
def query_non_agentic(self, user_query: str, stream: bool = None) -> Any:
|
||||
"""
|
||||
Non-agentic RAG mode: Simple retrieval + LLM response.
|
||||
|
||||
Args:
|
||||
user_query: The user's question
|
||||
stream: Whether to stream the response
|
||||
|
||||
Returns:
|
||||
The response (string or generator for streaming)
|
||||
"""
|
||||
if stream is None:
|
||||
stream = self.config.llm.stream
|
||||
|
||||
try:
|
||||
# Simple retrieval
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Non-agentic mode: searching for '{user_query}'")
|
||||
|
||||
search_results = self.kb_tools.knowledge_base_search(user_query)
|
||||
|
||||
if self.config.agent.verbose:
|
||||
logger.info(f"Non-agentic mode: found {len(search_results)} results")
|
||||
|
||||
# Build context from search results
|
||||
context_parts = []
|
||||
for i, result in enumerate(search_results[:5], 1): # Top 5 results
|
||||
context_parts.append(
|
||||
f"[Document {i}] (ID: {result['doc_id']}, Chunk: {result['chunk_id']})\n{result['text']}\n"
|
||||
)
|
||||
|
||||
if not context_parts:
|
||||
context = "No relevant information found in the knowledge base."
|
||||
else:
|
||||
context = "\n".join(context_parts)
|
||||
|
||||
# Build prompt
|
||||
system_prompt = """You are an assistant that answers questions based on provided context from a knowledge base.
|
||||
|
||||
IMPORTANT RULES:
|
||||
1. Only answer based on the provided context
|
||||
2. Include citations in format [Doc: document_id]
|
||||
3. If the context doesn't contain the answer, say so clearly
|
||||
4. Be accurate and don't make up information"""
|
||||
|
||||
user_prompt = f"""Context from knowledge base:
|
||||
{context}
|
||||
|
||||
User Question: {user_query}
|
||||
|
||||
Please answer the question based only on the provided context. Include citations."""
|
||||
|
||||
# Call LLM
|
||||
messages = [
|
||||
{"role": "system", "content": system_prompt},
|
||||
{"role": "user", "content": user_prompt}
|
||||
]
|
||||
|
||||
if stream:
|
||||
response_stream = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
|
||||
max_tokens=self.config.llm.max_tokens,
|
||||
stream=True
|
||||
)
|
||||
|
||||
def response_generator():
|
||||
for chunk in response_stream:
|
||||
if chunk.choices[0].delta.content:
|
||||
yield chunk.choices[0].delta.content
|
||||
|
||||
return response_generator()
|
||||
else:
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=messages,
|
||||
temperature=_reasoning_safe_temperature(self.model, self.config.llm.temperature),
|
||||
max_tokens=self.config.llm.max_tokens,
|
||||
stream=False
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in non-agentic query: {e}")
|
||||
error_msg = f"Error processing query: {str(e)}"
|
||||
if stream:
|
||||
return self._stream_response(error_msg)
|
||||
else:
|
||||
return error_msg
|
||||
|
||||
def clear_history(self):
|
||||
"""Clear conversation history"""
|
||||
self.conversation_history = []
|
||||
logger.info("Conversation history cleared")
|
||||
@@ -0,0 +1,283 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Canonical live plain-vs-contextual retrieval campaign (Experiment 3-10)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Sequence
|
||||
|
||||
import numpy as np
|
||||
from openai import OpenAI
|
||||
from rank_bm25 import BM25Okapi
|
||||
|
||||
HERE = Path(__file__).resolve().parent
|
||||
CHAPTER = HERE.parent
|
||||
sys.path.insert(0, str(CHAPTER))
|
||||
sys.path.insert(0, str(HERE))
|
||||
|
||||
from experiment_utils import ChatRecorder, sha256_file, write_campaign_evidence
|
||||
from compare_retrieval import tokenize
|
||||
|
||||
|
||||
ARK_ENDPOINT = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
|
||||
|
||||
class TransformerEncoder:
|
||||
def __init__(self, model_name: str, device: str):
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
self.torch = torch
|
||||
self.model_name = model_name
|
||||
self.device = device
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")
|
||||
self.model = AutoModel.from_pretrained(model_name).to(device).eval()
|
||||
self.revision = getattr(self.model.config, "_commit_hash", None)
|
||||
|
||||
def encode(self, texts: Sequence[str], *, query: bool, batch_size: int = 8) -> np.ndarray:
|
||||
prefix = "Instruct: Retrieve semantically relevant passages.\nQuery:" if query else ""
|
||||
vectors = []
|
||||
for start in range(0, len(texts), batch_size):
|
||||
batch = [prefix + text for text in texts[start : start + batch_size]]
|
||||
tokens = self.tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors="pt").to(self.device)
|
||||
with self.torch.no_grad():
|
||||
output = self.model(**tokens).last_hidden_state[:, -1].float()
|
||||
output = self.torch.nn.functional.normalize(output, p=2, dim=1)
|
||||
vectors.append(output.cpu().numpy())
|
||||
return np.concatenate(vectors).astype("float32")
|
||||
|
||||
|
||||
def load_chunks(path: Path) -> List[Dict[str, Any]]:
|
||||
store = json.loads(path.read_text(encoding="utf-8"))
|
||||
rows = []
|
||||
for chunk_id, entry in store.items():
|
||||
if "_chunk_" not in chunk_id:
|
||||
continue
|
||||
meta = entry.get("metadata") or {}
|
||||
rows.append(
|
||||
{
|
||||
"chunk_id": chunk_id,
|
||||
"doc_title": meta.get("doc_title") or chunk_id.split("_chunk_")[0],
|
||||
"plain": meta.get("original_text") or entry.get("content", ""),
|
||||
}
|
||||
)
|
||||
return sorted(rows, key=lambda row: row["chunk_id"])
|
||||
|
||||
|
||||
def source_documents(chunks: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
|
||||
laws = CHAPTER / "agentic-rag" / "laws"
|
||||
output = {}
|
||||
for title in sorted({row["doc_title"] for row in chunks}):
|
||||
candidates = [path for path in laws.rglob("*.md") if path.stem == title]
|
||||
if len(candidates) != 1:
|
||||
raise RuntimeError(f"expected one official bundled source for {title!r}, found {len(candidates)}")
|
||||
path = candidates[0]
|
||||
output[title] = {"path": path, "text": path.read_text(encoding="utf-8")}
|
||||
return output
|
||||
|
||||
|
||||
def prefix_one(args: argparse.Namespace, chunk: Dict[str, Any], source: Dict[str, Any]):
|
||||
client = OpenAI(api_key=os.environ["ARK_API_KEY"], base_url=args.endpoint, timeout=args.timeout, max_retries=3)
|
||||
recorder = ChatRecorder(client, "ark", args.endpoint)
|
||||
response = recorder.create(
|
||||
purpose=f"3-10 live contextual prefix {chunk['chunk_id']}",
|
||||
model=args.context_model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"为目标文本块生成简短的中文检索前缀。前缀必须说明该块来自哪份文档、所属章节/条款、"
|
||||
"主体与主题,使孤立文本能被准确检索。不得添加源文没有的事实。只输出前缀,不要解释。"
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"完整源文档:\n<document>\n{source['text']}\n</document>\n\n目标文本块:\n<chunk>\n{chunk['plain']}\n</chunk>",
|
||||
},
|
||||
],
|
||||
temperature=0,
|
||||
seed=args.seed,
|
||||
max_tokens=220,
|
||||
)
|
||||
prefix = (response.choices[0].message.content or "").strip()
|
||||
return {**chunk, "prefix": prefix, "contextual": f"{prefix}\n\n{chunk['plain']}"}, recorder.calls
|
||||
|
||||
|
||||
def rankings_bm25(texts: List[str], queries: List[str]) -> List[List[int]]:
|
||||
index = BM25Okapi([tokenize(text) for text in texts])
|
||||
return [np.argsort(-index.get_scores(tokenize(query))).tolist() for query in queries]
|
||||
|
||||
|
||||
def rankings_dense(vectors: np.ndarray, query_vectors: np.ndarray) -> List[List[int]]:
|
||||
return [np.argsort(-(query @ vectors.T)).tolist() for query in query_vectors]
|
||||
|
||||
|
||||
def rrf(a: List[int], b: List[int], constant: int = 60) -> List[int]:
|
||||
scores: Dict[int, float] = {}
|
||||
for ranking in (a, b):
|
||||
for rank, item in enumerate(ranking, start=1):
|
||||
scores[item] = scores.get(item, 0.0) + 1.0 / (constant + rank)
|
||||
return sorted(scores, key=lambda item: scores[item], reverse=True)
|
||||
|
||||
|
||||
def metrics(rankings: List[List[int]], queries: List[Dict[str, Any]], id_to_pos: Dict[str, int]) -> Dict[str, Any]:
|
||||
per_query = []
|
||||
reciprocal = []
|
||||
for query, ranking in zip(queries, rankings):
|
||||
gold = id_to_pos[query["gold_chunk_id"]]
|
||||
rank = ranking.index(gold) + 1 if gold in ranking else None
|
||||
reciprocal.append(1.0 / rank if rank else 0.0)
|
||||
per_query.append(
|
||||
{
|
||||
"id": query["id"],
|
||||
"query": query["query"],
|
||||
"gold_chunk_id": query["gold_chunk_id"],
|
||||
"rank": rank,
|
||||
"top5_chunk_ids": ranking[:5],
|
||||
}
|
||||
)
|
||||
return {
|
||||
"n": len(queries),
|
||||
"recall_at_k": {str(k): statistics.mean(1.0 if row["rank"] and row["rank"] <= k else 0.0 for row in per_query) for k in (1, 3, 5)},
|
||||
"mrr": statistics.mean(reciprocal),
|
||||
"per_query": per_query,
|
||||
}
|
||||
|
||||
|
||||
def token_usage(receipts: List[Dict[str, Any]]) -> Dict[str, int]:
|
||||
totals = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
||||
for call in receipts:
|
||||
usage = call.get("usage") or {}
|
||||
for key in totals:
|
||||
totals[key] += int(usage.get(key) or 0)
|
||||
return totals
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--context-model", default=os.getenv("ARK_MODEL", "doubao-seed-1-6-250615"))
|
||||
parser.add_argument("--embedding-model", default="Qwen/Qwen3-Embedding-0.6B")
|
||||
parser.add_argument("--device", default="cpu")
|
||||
parser.add_argument("--endpoint", default=ARK_ENDPOINT)
|
||||
parser.add_argument("--workers", type=int, default=4)
|
||||
parser.add_argument("--seed", type=int, default=37)
|
||||
parser.add_argument("--timeout", type=float, default=180)
|
||||
parser.add_argument("--input-price-per-million-usd", type=float, default=0.11)
|
||||
parser.add_argument("--output-price-per-million-usd", type=float, default=1.10)
|
||||
args = parser.parse_args()
|
||||
if not os.getenv("ARK_API_KEY"):
|
||||
raise RuntimeError("ARK_API_KEY is required")
|
||||
|
||||
corpus_path = HERE / "document_store.json"
|
||||
eval_path = HERE / "evaluation" / "retrieval_eval.json"
|
||||
chunks = load_chunks(corpus_path)
|
||||
docs = source_documents(chunks)
|
||||
receipts: List[Dict[str, Any]] = []
|
||||
contextual: List[Dict[str, Any]] = []
|
||||
errors = []
|
||||
prefix_start = time.perf_counter()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
|
||||
futures = {pool.submit(prefix_one, args, chunk, docs[chunk["doc_title"]]): chunk["chunk_id"] for chunk in chunks}
|
||||
for future in concurrent.futures.as_completed(futures):
|
||||
chunk_id = futures[future]
|
||||
try:
|
||||
row, calls = future.result()
|
||||
contextual.append(row)
|
||||
receipts.extend(calls)
|
||||
print(f"prefix {chunk_id} ({len(contextual)}/{len(chunks)})", flush=True)
|
||||
except Exception as exc:
|
||||
errors.append({"chunk_id": chunk_id, "type": type(exc).__name__, "error": str(exc)})
|
||||
prefix_ms = (time.perf_counter() - prefix_start) * 1000
|
||||
contextual.sort(key=lambda row: row["chunk_id"])
|
||||
|
||||
eval_data = json.loads(eval_path.read_text(encoding="utf-8"))
|
||||
queries = eval_data["queries"]
|
||||
query_texts = [row["query"] for row in queries]
|
||||
ids = [row["chunk_id"] for row in contextual]
|
||||
id_to_pos = {chunk_id: pos for pos, chunk_id in enumerate(ids)}
|
||||
methods: Dict[str, Dict[str, Any]] = {}
|
||||
embedding_ms = 0.0
|
||||
encoder = None
|
||||
if len(contextual) == len(chunks) and not errors:
|
||||
plain_texts = [row["plain"] for row in contextual]
|
||||
contextual_texts = [row["contextual"] for row in contextual]
|
||||
plain_bm25 = rankings_bm25(plain_texts, query_texts)
|
||||
contextual_bm25 = rankings_bm25(contextual_texts, query_texts)
|
||||
encoder = TransformerEncoder(args.embedding_model, args.device)
|
||||
started = time.perf_counter()
|
||||
plain_vectors = encoder.encode(plain_texts, query=False)
|
||||
contextual_vectors = encoder.encode(contextual_texts, query=False)
|
||||
query_vectors = encoder.encode(query_texts, query=True)
|
||||
embedding_ms = (time.perf_counter() - started) * 1000
|
||||
plain_dense = rankings_dense(plain_vectors, query_vectors)
|
||||
contextual_dense = rankings_dense(contextual_vectors, query_vectors)
|
||||
ranking_sets = {
|
||||
"plain_bm25": plain_bm25,
|
||||
"contextual_bm25": contextual_bm25,
|
||||
"plain_dense": plain_dense,
|
||||
"contextual_dense": contextual_dense,
|
||||
"plain_hybrid": [rrf(a, b) for a, b in zip(plain_bm25, plain_dense)],
|
||||
"contextual_hybrid": [rrf(a, b) for a, b in zip(contextual_bm25, contextual_dense)],
|
||||
}
|
||||
for name, ranking in ranking_sets.items():
|
||||
result = metrics(ranking, queries, id_to_pos)
|
||||
for row in result["per_query"]:
|
||||
row["top5_chunk_ids"] = [ids[pos] for pos in row["top5_chunk_ids"]]
|
||||
methods[name] = result
|
||||
|
||||
tokens = token_usage(receipts)
|
||||
estimated_cost = tokens["prompt_tokens"] / 1_000_000 * args.input_price_per_million_usd + tokens["completion_tokens"] / 1_000_000 * args.output_price_per_million_usd
|
||||
acceptance = {
|
||||
"live_prefix_for_every_chunk": len(contextual) == len(chunks) and all(row["prefix"] for row in contextual),
|
||||
"full_source_document_and_target_chunk_in_requests": len(receipts) == len(chunks) and all("<document>" in json.dumps(call.get("request", {}), ensure_ascii=False) and "<chunk>" in json.dumps(call.get("request", {}), ensure_ascii=False) for call in receipts),
|
||||
"same_chunks_and_queries": bool(methods) and all(result["n"] == len(queries) for result in methods.values()),
|
||||
"plain_contextual_bm25_dense_hybrid": set(methods) == {"plain_bm25", "contextual_bm25", "plain_dense", "contextual_dense", "plain_hybrid", "contextual_hybrid"},
|
||||
"recall_and_mrr_measured": bool(methods) and all("mrr" in result and set(result["recall_at_k"]) == {"1", "3", "5"} for result in methods.values()),
|
||||
"real_dense_model": bool(encoder and encoder.revision),
|
||||
"index_usage_and_cost_measured": tokens["total_tokens"] > 0 and estimated_cost >= 0,
|
||||
"raw_request_response_receipts": len(receipts) == len(chunks) and all("request" in call and "response" in call for call in receipts),
|
||||
"all_calls_succeeded": not errors,
|
||||
}
|
||||
acceptance["passed"] = all(acceptance.values())
|
||||
evidence = {
|
||||
"status": "passed" if acceptance["passed"] else ("partial" if contextual else "blocked"),
|
||||
"configuration": vars(args) | {"embedding_revision": getattr(encoder, "revision", None)},
|
||||
"scope": {"documents": len(docs), "chunks": len(chunks), "queries": len(queries)},
|
||||
"acceptance": acceptance,
|
||||
"summary": {
|
||||
"methods": {name: {key: value for key, value in result.items() if key != "per_query"} for name, result in methods.items()},
|
||||
"index_time": {
|
||||
"context_generation_ms": round(prefix_ms, 3),
|
||||
"embedding_ms": round(embedding_ms, 3),
|
||||
"usage": tokens,
|
||||
"estimated_cost_usd": round(estimated_cost, 6),
|
||||
"pricing_assumption": {"input_per_million_usd": args.input_price_per_million_usd, "output_per_million_usd": args.output_price_per_million_usd},
|
||||
},
|
||||
"errors": len(errors),
|
||||
},
|
||||
"errors": errors,
|
||||
"source_documents": {title: {"path": str(data["path"]), "sha256": sha256_file(data["path"])} for title, data in docs.items()},
|
||||
"chunks": contextual,
|
||||
"results": methods,
|
||||
}
|
||||
manifest = write_campaign_evidence(
|
||||
HERE,
|
||||
"3-10",
|
||||
evidence,
|
||||
receipts,
|
||||
input_paths=[HERE / "campaign.py", HERE / "compare_retrieval.py", corpus_path, eval_path, *[data["path"] for data in docs.values()]],
|
||||
)
|
||||
print(json.dumps(manifest["summary"], ensure_ascii=False, indent=2))
|
||||
print(f"Canonical evidence: {HERE / 'validation' / 'latest.json'}")
|
||||
return 0 if acceptance["passed"] else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,423 @@
|
||||
"""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 = []
|
||||
# Step to the end of the list: re.split with a capturing group yields
|
||||
# [text, delim, text, delim, ..., trailing_text], so stopping at
|
||||
# len(sentences) - 1 dropped the trailing fragment whenever the text
|
||||
# did not end in terminal punctuation (and returned [] for text with
|
||||
# none at all). The strip-and-filter below still discards the empty
|
||||
# tail that re.split produces when the text does end in punctuation.
|
||||
for i in range(0, len(sentences), 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"]
|
||||
}
|
||||
}, timeout=30
|
||||
)
|
||||
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, timeout=30
|
||||
)
|
||||
else:
|
||||
# Create new document
|
||||
response = requests.post(
|
||||
f"{self.kb_config.dify_base_url}/documents",
|
||||
headers=headers,
|
||||
json=payload, timeout=30
|
||||
)
|
||||
|
||||
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()
|
||||
@@ -0,0 +1,344 @@
|
||||
#!/usr/bin/env python3
|
||||
"""上下文感知检索对比评测(实验 3-10)
|
||||
|
||||
本脚本用可控的对比实验量化“上下文感知检索”相较传统分块的检索召回提升:
|
||||
同一批文本块分别以两种方式建立 BM25 索引——
|
||||
|
||||
* 无上下文(plain) :只索引原始文本块 metadata.original_text
|
||||
* 有上下文(contextual):索引 LLM 生成的前缀 + 原始文本块(content 字段)
|
||||
|
||||
然后在同一评测集上比较 recall@k(命中率:前 k 个结果中是否含有相关文本块)。
|
||||
这正是 Anthropic “Contextual Retrieval” 的核心主张:为文本块补上上下文前缀,
|
||||
能同时增强 BM25(稀疏)与向量(稠密)检索的召回率。
|
||||
|
||||
BM25 检索完全离线,无需任何 API 或检索服务;embedding / hybrid 方法需要
|
||||
调用 embedding API(见 --method 说明)。
|
||||
|
||||
用法示例:
|
||||
python compare_retrieval.py # 用默认评测集跑对比表
|
||||
python compare_retrieval.py --query "国家主席有哪些职权?" # 单条查询并排对比
|
||||
python compare_retrieval.py --mode plain # 只看无上下文基线
|
||||
python compare_retrieval.py --output result.json # 另存机器可读结果
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from rank_bm25 import BM25Okapi
|
||||
|
||||
try:
|
||||
import jieba
|
||||
if hasattr(jieba, "setLogLevel"):
|
||||
jieba.setLogLevel(60) # 关闭 jieba 的加载日志
|
||||
_HAS_JIEBA = True
|
||||
except Exception: # pragma: no cover - jieba 一般随 requirements 安装
|
||||
_HAS_JIEBA = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 分词:中文没有空格,直接 .split() 会把整段当成一个 token,BM25 完全失效。
|
||||
# 默认用 jieba 分词;--no-jieba 时退化为字符二元组(bigram),同样可离线运行。
|
||||
# ---------------------------------------------------------------------------
|
||||
def tokenize(text: str, use_jieba: bool = True) -> List[str]:
|
||||
"""把文本切成 token 列表,供 BM25 使用。"""
|
||||
text = (text or "").lower()
|
||||
if use_jieba and _HAS_JIEBA:
|
||||
return [t for t in jieba.cut(text) if t.strip()]
|
||||
# 退化方案:中文字符二元组 + 连续 ASCII 词
|
||||
tokens: List[str] = []
|
||||
buf = ""
|
||||
chars = list(text)
|
||||
for ch in chars:
|
||||
if ch.isascii() and (ch.isalnum()):
|
||||
buf += ch
|
||||
continue
|
||||
if buf:
|
||||
tokens.append(buf)
|
||||
buf = ""
|
||||
if not ch.isspace():
|
||||
tokens.append(ch)
|
||||
if buf:
|
||||
tokens.append(buf)
|
||||
# 追加中文 bigram,提升匹配粒度
|
||||
cjk = [c for c in text if "一" <= c <= "鿿"]
|
||||
tokens.extend(cjk[i] + cjk[i + 1] for i in range(len(cjk) - 1))
|
||||
return tokens
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 语料加载
|
||||
# ---------------------------------------------------------------------------
|
||||
def load_corpus(path: str) -> List[Dict]:
|
||||
"""从 document_store.json 载入分块,返回 [{chunk_id, contextual, plain, context}]。
|
||||
|
||||
每个分块的 content 字段是“上下文前缀 + 原始文本”,metadata.original_text
|
||||
是不带上下文的原始文本,正好用于两种索引方式的对照。
|
||||
"""
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
store = json.load(f)
|
||||
|
||||
chunks: List[Dict] = []
|
||||
for chunk_id, entry in store.items():
|
||||
if "_chunk_" not in chunk_id:
|
||||
continue # 跳过整篇文档条目
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
meta = entry.get("metadata", {}) or {}
|
||||
contextual_text = entry.get("content", "") or ""
|
||||
plain_text = meta.get("original_text") or contextual_text
|
||||
# 上下文前缀 = contextual 去掉结尾的 original_text
|
||||
context = contextual_text
|
||||
if plain_text and contextual_text.endswith(plain_text):
|
||||
context = contextual_text[: len(contextual_text) - len(plain_text)].strip()
|
||||
chunks.append({
|
||||
"chunk_id": chunk_id,
|
||||
"contextual": contextual_text,
|
||||
"plain": plain_text,
|
||||
"context": context,
|
||||
})
|
||||
return chunks
|
||||
|
||||
|
||||
def load_eval(path: str) -> List[Dict]:
|
||||
"""载入评测集,返回 [{id, query, gold_chunk_id, ...}]。"""
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return data.get("queries", data if isinstance(data, list) else [])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# BM25 检索器
|
||||
# ---------------------------------------------------------------------------
|
||||
class BM25Retriever:
|
||||
"""对给定文本字段建立 BM25 索引的简单检索器。"""
|
||||
|
||||
def __init__(self, chunks: List[Dict], field: str, use_jieba: bool = True):
|
||||
self.chunk_ids = [c["chunk_id"] for c in chunks]
|
||||
self.use_jieba = use_jieba
|
||||
corpus_tokens = [tokenize(c[field], use_jieba) for c in chunks]
|
||||
self.index = BM25Okapi(corpus_tokens)
|
||||
|
||||
def rank(self, query: str) -> List[str]:
|
||||
"""返回按相关性从高到低排序的 chunk_id 列表。"""
|
||||
scores = self.index.get_scores(tokenize(query, self.use_jieba))
|
||||
order = np.argsort(scores)[::-1]
|
||||
return [self.chunk_ids[i] for i in order]
|
||||
|
||||
def scored(self, query: str, top_k: int) -> List[Dict]:
|
||||
"""返回前 top_k 个结果及其分数。"""
|
||||
scores = self.index.get_scores(tokenize(query, self.use_jieba))
|
||||
order = np.argsort(scores)[::-1][:top_k]
|
||||
return [{"chunk_id": self.chunk_ids[i], "score": float(scores[i])} for i in order]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 评测
|
||||
# ---------------------------------------------------------------------------
|
||||
def recall_at_k(retriever: BM25Retriever, queries: List[Dict], ks: List[int]) -> Dict:
|
||||
"""计算一批查询在各 k 值下的 recall@k(命中率)。"""
|
||||
per_query = []
|
||||
hits = {k: 0 for k in ks}
|
||||
for q in queries:
|
||||
ranking = retriever.rank(q["query"])
|
||||
gold = q["gold_chunk_id"]
|
||||
rank_pos = ranking.index(gold) + 1 if gold in ranking else None
|
||||
row = {"id": q.get("id"), "query": q["query"], "gold": gold, "rank": rank_pos}
|
||||
for k in ks:
|
||||
hit = rank_pos is not None and rank_pos <= k
|
||||
row[f"hit@{k}"] = hit
|
||||
if hit:
|
||||
hits[k] += 1
|
||||
per_query.append(row)
|
||||
n = len(queries)
|
||||
recall = {k: (hits[k] / n if n else 0.0) for k in ks}
|
||||
return {"recall": recall, "per_query": per_query, "n": n}
|
||||
|
||||
|
||||
def print_comparison_table(plain: Optional[Dict], contextual: Optional[Dict], ks: List[int]):
|
||||
"""打印 recall@k 对比表。"""
|
||||
print("\n" + "=" * 68)
|
||||
print("检索召回对比:无上下文分块 vs. 上下文感知检索(BM25)")
|
||||
print("=" * 68)
|
||||
header = " k | " + " | ".join(f"{'无上下文':>10}" if False else f"recall@{k:<3}" for k in ks)
|
||||
# 逐行打印每个方法
|
||||
col_w = 12
|
||||
line = f"{'方法':<16}" + "".join(f"recall@{k}".rjust(col_w) for k in ks)
|
||||
print(line)
|
||||
print("-" * len(line))
|
||||
if plain:
|
||||
print(f"{'无上下文 (plain)':<16}" + "".join(f"{plain['recall'][k]*100:>10.1f}%" for k in ks))
|
||||
if contextual:
|
||||
print(f"{'有上下文 (ctx)':<16}" + "".join(f"{contextual['recall'][k]*100:>10.1f}%" for k in ks))
|
||||
if plain and contextual:
|
||||
print("-" * len(line))
|
||||
deltas = []
|
||||
for k in ks:
|
||||
d = (contextual["recall"][k] - plain["recall"][k]) * 100
|
||||
deltas.append(f"{d:>+9.1f}pp")
|
||||
print(f"{'提升 (Δpp)':<16}" + "".join(s.rjust(col_w) for s in deltas))
|
||||
# 检索失败率下降(对应书中“1 - recall@k”口径)
|
||||
print("-" * len(line))
|
||||
fails = []
|
||||
for k in ks:
|
||||
p_fail = 1 - plain["recall"][k]
|
||||
c_fail = 1 - contextual["recall"][k]
|
||||
if p_fail > 0:
|
||||
red = (p_fail - c_fail) / p_fail * 100
|
||||
fails.append(f"{red:>9.0f}%")
|
||||
else:
|
||||
fails.append(f"{'-':>10}")
|
||||
print(f"{'失败率下降':<16}" + "".join(s.rjust(col_w) for s in fails))
|
||||
print("=" * 68)
|
||||
|
||||
|
||||
def print_per_query(result: Dict, label: str):
|
||||
print(f"\n[{label}] 每条查询命中排名(rank=gold 文本块在结果中的名次,— 表示未召回)")
|
||||
for row in result["per_query"]:
|
||||
print(f" {row['id']} rank={str(row['rank']):>3} gold={row['gold']:<28} {row['query'][:32]}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 单条查询并排对比
|
||||
# ---------------------------------------------------------------------------
|
||||
def single_query_compare(chunks: List[Dict], query: str, top_k: int, use_jieba: bool,
|
||||
mode: str):
|
||||
id2chunk = {c["chunk_id"]: c for c in chunks}
|
||||
|
||||
def show(field_label, field):
|
||||
retr = BM25Retriever(chunks, field, use_jieba)
|
||||
print(f"\n[{field_label}] Top-{top_k}")
|
||||
print("-" * 60)
|
||||
for i, r in enumerate(retr.scored(query, top_k), 1):
|
||||
c = id2chunk[r["chunk_id"]]
|
||||
snippet = c["plain"].replace("<!-- FORCE BREAK -->", "").replace("\n", " ").strip()[:48]
|
||||
ctx = c["context"].replace("\n", " ").strip()[:40]
|
||||
print(f" {i}. score={r['score']:6.2f} {r['chunk_id']}")
|
||||
if field == "contextual" and ctx:
|
||||
print(f" 上下文前缀: {ctx}")
|
||||
print(f" 原文: {snippet}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(f"查询: {query}")
|
||||
print("=" * 60)
|
||||
if mode in ("plain", "both"):
|
||||
show("无上下文 (plain)", "plain")
|
||||
if mode in ("contextual", "both"):
|
||||
show("有上下文 (contextual)", "contextual")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 可选:embedding / hybrid(需要 API)
|
||||
# ---------------------------------------------------------------------------
|
||||
def embedding_unavailable_notice(method: str):
|
||||
print(f"\n[提示] --method {method} 需要调用 embedding API(稠密向量),无法离线运行。")
|
||||
print(" 请在 .env 中配置 OPENAI_API_KEY / SILICONFLOW_API_KEY 等,")
|
||||
print(" 并使用 contextual_tools.ContextualKnowledgeBaseTools 的 embedding/hybrid 检索。")
|
||||
print(" 本脚本的默认 --method bm25 已可完整复现书中“上下文增强 BM25”的召回提升结论。")
|
||||
|
||||
|
||||
def build_arg_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(
|
||||
description="上下文感知检索对比评测:量化上下文前缀对检索召回(recall@k)的提升(实验 3-10)",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="示例:\n"
|
||||
" python compare_retrieval.py\n"
|
||||
" python compare_retrieval.py --query \"国家主席有哪些职权?\" --top-k 5\n"
|
||||
" python compare_retrieval.py --mode both --k 1 3 5 --output result.json",
|
||||
)
|
||||
p.add_argument("--corpus", default="document_store.json",
|
||||
help="语料文件(含 content 与 metadata.original_text 的分块存储),默认 document_store.json")
|
||||
p.add_argument("--eval", dest="eval_path", default="evaluation/retrieval_eval.json",
|
||||
help="评测集(query + gold_chunk_id),默认 evaluation/retrieval_eval.json")
|
||||
p.add_argument("--query", default=None,
|
||||
help="临时单条查询:并排展示无上下文/有上下文的 Top-K 检索结果(不跑整个评测集)")
|
||||
p.add_argument("--mode", choices=["plain", "contextual", "both"], default="both",
|
||||
help="对比哪种索引:plain=仅无上下文,contextual=仅有上下文,both=两者对比(默认)")
|
||||
p.add_argument("--method", choices=["bm25", "embedding", "hybrid"], default="bm25",
|
||||
help="检索方法:bm25(离线,默认);embedding/hybrid 需 embedding API")
|
||||
p.add_argument("--k", nargs="+", type=int, default=[1, 3, 5],
|
||||
help="评测的 k 值列表(recall@k),默认 1 3 5")
|
||||
p.add_argument("--top-k", type=int, default=5,
|
||||
help="--query 单查询模式下每种方法展示的结果条数,默认 5")
|
||||
p.add_argument("--model", default=None,
|
||||
help="embedding 模型名(仅 --method embedding/hybrid 时生效)")
|
||||
p.add_argument("--no-jieba", action="store_true",
|
||||
help="禁用 jieba 分词,改用字符二元组分词(无需 jieba 依赖)")
|
||||
p.add_argument("--output", default=None,
|
||||
help="将机器可读的评测结果写入该 JSON 文件")
|
||||
p.add_argument("--per-query", action="store_true",
|
||||
help="额外打印每条查询的命中排名明细")
|
||||
return p
|
||||
|
||||
|
||||
def main():
|
||||
args = build_arg_parser().parse_args()
|
||||
use_jieba = not args.no_jieba
|
||||
|
||||
corpus_path = Path(args.corpus)
|
||||
if not corpus_path.exists():
|
||||
print(f"[错误] 找不到语料文件: {corpus_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
chunks = load_corpus(str(corpus_path))
|
||||
if not chunks:
|
||||
print(f"[错误] 语料中没有可用分块(缺少 *_chunk_* 条目): {corpus_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
print(f"已加载 {len(chunks)} 个文本块 | 分词: {'jieba' if (use_jieba and _HAS_JIEBA) else '字符bigram'} "
|
||||
f"| 检索方法: {args.method}")
|
||||
|
||||
if args.method in ("embedding", "hybrid"):
|
||||
embedding_unavailable_notice(args.method)
|
||||
# 仍继续用 BM25 给出可运行的离线结果
|
||||
print(" 以下改用 BM25 给出离线对照结果。\n")
|
||||
|
||||
# 单条查询模式
|
||||
if args.query:
|
||||
single_query_compare(chunks, args.query, args.top_k, use_jieba, args.mode)
|
||||
return
|
||||
|
||||
# 评测集模式
|
||||
eval_path = Path(args.eval_path)
|
||||
if not eval_path.exists():
|
||||
print(f"[错误] 找不到评测集: {eval_path}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
queries = load_eval(str(eval_path))
|
||||
ks = sorted(set(args.k))
|
||||
|
||||
plain_res = contextual_res = None
|
||||
if args.mode in ("plain", "both"):
|
||||
plain_res = recall_at_k(BM25Retriever(chunks, "plain", use_jieba), queries, ks)
|
||||
if args.mode in ("contextual", "both"):
|
||||
contextual_res = recall_at_k(BM25Retriever(chunks, "contextual", use_jieba), queries, ks)
|
||||
|
||||
print(f"评测集: {eval_path} 共 {len(queries)} 条查询")
|
||||
print_comparison_table(plain_res, contextual_res, ks)
|
||||
|
||||
if args.per_query:
|
||||
if plain_res:
|
||||
print_per_query(plain_res, "无上下文 plain")
|
||||
if contextual_res:
|
||||
print_per_query(contextual_res, "有上下文 contextual")
|
||||
|
||||
if args.output:
|
||||
out = {
|
||||
"corpus": str(corpus_path),
|
||||
"eval": str(eval_path),
|
||||
"num_chunks": len(chunks),
|
||||
"num_queries": len(queries),
|
||||
"tokenizer": "jieba" if (use_jieba and _HAS_JIEBA) else "char-bigram",
|
||||
"k": ks,
|
||||
"plain": plain_res,
|
||||
"contextual": contextual_res,
|
||||
}
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
json.dump(out, f, ensure_ascii=False, indent=2)
|
||||
print(f"\n结果已写入 {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,259 @@
|
||||
"""Configuration for Agentic RAG System"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Any
|
||||
from enum import Enum
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
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"""
|
||||
DASHSCOPE = "dashscope" # Alibaba Cloud Model Studio / Bailian (Qwen)
|
||||
SILICONFLOW = "siliconflow"
|
||||
DOUBAO = "doubao"
|
||||
KIMI = "kimi"
|
||||
MOONSHOT = "moonshot"
|
||||
OPENROUTER = "openrouter"
|
||||
OPENAI = "openai"
|
||||
GROQ = "groq"
|
||||
TOGETHER = "together"
|
||||
DEEPSEEK = "deepseek"
|
||||
|
||||
|
||||
class KnowledgeBaseType(str, Enum):
|
||||
"""Knowledge base backend types"""
|
||||
LOCAL = "local" # Local retrieval pipeline
|
||||
DIFY = "dify" # Dify knowledge base API
|
||||
RAPTOR = "raptor" # RAPTOR tree-based index
|
||||
GRAPHRAG = "graphrag" # GraphRAG graph-based index
|
||||
|
||||
|
||||
@dataclass
|
||||
class LLMConfig:
|
||||
"""LLM configuration"""
|
||||
provider: str = "kimi" # Default provider
|
||||
model: Optional[str] = None # Will use provider defaults if not specified
|
||||
api_key: Optional[str] = None # Will read from env if not provided
|
||||
temperature: float = 0.7
|
||||
max_tokens: int = 1024
|
||||
stream: bool = True
|
||||
|
||||
# Provider-specific defaults
|
||||
PROVIDER_DEFAULTS = {
|
||||
"dashscope": {
|
||||
"model": "qwen3.7-plus",
|
||||
"base_url": os.getenv(
|
||||
"DASHSCOPE_BASE_URL",
|
||||
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
),
|
||||
},
|
||||
"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 = {
|
||||
"dashscope": "DASHSCOPE_API_KEY",
|
||||
"qwen": "DASHSCOPE_API_KEY",
|
||||
"bailian": "DASHSCOPE_API_KEY",
|
||||
"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()
|
||||
provider_lower = {"qwen": "dashscope", "bailian": "dashscope"}.get(
|
||||
provider_lower, provider_lower
|
||||
)
|
||||
defaults = self.PROVIDER_DEFAULTS.get(provider_lower, {})
|
||||
|
||||
# Get API key
|
||||
api_key = self.api_key or self.get_api_key(provider_lower)
|
||||
|
||||
# Universal OpenRouter fallback: primary provider key absent but
|
||||
# OPENROUTER_API_KEY present -> route through OpenRouter. Additionally,
|
||||
# gpt-5.x (incl. gpt-5.6*) needs OpenAI org-verification on the direct
|
||||
# API, so prefer OpenRouter for those ids whenever an OR key is present.
|
||||
model_name = self.model or defaults.get("model")
|
||||
openrouter_key = os.getenv("OPENROUTER_API_KEY")
|
||||
prefer_openrouter = bool(openrouter_key) and str(model_name or "").lower().startswith("gpt-5")
|
||||
if (not api_key or prefer_openrouter) and provider_lower != "openrouter" and openrouter_key:
|
||||
model = _openrouter_model_id(model_name)
|
||||
return {
|
||||
"api_key": openrouter_key,
|
||||
"base_url": "https://openrouter.ai/api/v1",
|
||||
}, model
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
f"API key required for provider '{provider_lower}'. Set the "
|
||||
f"provider's key (e.g. MOONSHOT_API_KEY / OPENAI_API_KEY) or "
|
||||
f"OPENROUTER_API_KEY to use the OpenRouter fallback."
|
||||
)
|
||||
|
||||
# Build config
|
||||
config = {
|
||||
"api_key": api_key,
|
||||
"model": self.model or defaults.get("model")
|
||||
}
|
||||
|
||||
# Add base_url if not OpenAI
|
||||
if "base_url" in defaults:
|
||||
config["base_url"] = defaults["base_url"]
|
||||
|
||||
return config, config.pop("model")
|
||||
|
||||
|
||||
@dataclass
|
||||
class KnowledgeBaseConfig:
|
||||
"""Knowledge base configuration"""
|
||||
type: KnowledgeBaseType = KnowledgeBaseType.LOCAL
|
||||
|
||||
# Local retrieval pipeline config
|
||||
local_base_url: str = "http://localhost:4242"
|
||||
local_top_k: int = 3
|
||||
|
||||
# Dify config
|
||||
dify_api_key: Optional[str] = field(default_factory=lambda: os.getenv("DIFY_API_KEY"))
|
||||
dify_base_url: str = "https://api.dify.ai/v1"
|
||||
dify_dataset_id: Optional[str] = None
|
||||
dify_top_k: int = 10
|
||||
|
||||
# RAPTOR tree-based index config
|
||||
raptor_base_url: str = "http://localhost:4242"
|
||||
raptor_top_k: int = 10
|
||||
raptor_search_levels: bool = True # Search across multiple tree levels
|
||||
|
||||
# GraphRAG graph-based index config
|
||||
graphrag_base_url: str = "http://localhost:4242"
|
||||
graphrag_top_k: int = 10
|
||||
graphrag_search_type: str = "hybrid" # entity, community, or hybrid
|
||||
|
||||
# Document storage
|
||||
document_store_path: str = "document_store.json"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkingConfig:
|
||||
"""Document chunking configuration"""
|
||||
chunk_size: int = 2048 # Characters per chunk
|
||||
max_chunk_size: int = 1024 # Max size when respecting paragraph boundaries
|
||||
chunk_overlap: int = 200 # Overlap between chunks
|
||||
respect_paragraph_boundary: bool = True
|
||||
min_chunk_size: int = 100 # Minimum chunk size
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""Agent configuration"""
|
||||
max_iterations: int = 10 # Max reasoning iterations
|
||||
enable_reasoning_trace: bool = True
|
||||
enable_citations: bool = True
|
||||
strict_knowledge_base: bool = True # Only answer from knowledge base
|
||||
conversation_history_limit: int = 20 # Max conversation turns to keep
|
||||
verbose: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class EvaluationConfig:
|
||||
"""Evaluation configuration"""
|
||||
dataset_path: str = "evaluation/legal_qa_dataset.json"
|
||||
results_path: str = "evaluation/results"
|
||||
metrics: list = field(default_factory=lambda: ["accuracy", "relevance", "citation_quality"])
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Main configuration"""
|
||||
llm: LLMConfig = field(default_factory=LLMConfig)
|
||||
knowledge_base: KnowledgeBaseConfig = field(default_factory=KnowledgeBaseConfig)
|
||||
chunking: ChunkingConfig = field(default_factory=ChunkingConfig)
|
||||
agent: AgentConfig = field(default_factory=AgentConfig)
|
||||
evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls) -> "Config":
|
||||
"""Create config from environment variables"""
|
||||
config = cls()
|
||||
|
||||
# Override from env
|
||||
if provider := os.getenv("LLM_PROVIDER"):
|
||||
config.llm.provider = provider
|
||||
if model := os.getenv("LLM_MODEL"):
|
||||
config.llm.model = model
|
||||
if kb_type := os.getenv("KB_TYPE"):
|
||||
config.knowledge_base.type = KnowledgeBaseType(kb_type.lower())
|
||||
|
||||
return config
|
||||
@@ -0,0 +1,356 @@
|
||||
"""Contextual Agentic RAG System
|
||||
|
||||
This module extends the base AgenticRAG to use contextual retrieval,
|
||||
demonstrating the improved answer quality from better retrieval.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional, Generator
|
||||
from datetime import datetime
|
||||
|
||||
from agent import AgenticRAG, Message
|
||||
from config import Config
|
||||
from contextual_tools import ContextualKnowledgeBaseTools, ContextualSearchResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ContextualAgenticRAG(AgenticRAG):
|
||||
"""
|
||||
Enhanced Agentic RAG with contextual retrieval support.
|
||||
|
||||
Educational Features:
|
||||
- Shows how better retrieval leads to better answers
|
||||
- Logs retrieval quality metrics
|
||||
- Compares contextual vs non-contextual results
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config: Optional[Config] = None,
|
||||
kb_tools: Optional[ContextualKnowledgeBaseTools] = None,
|
||||
use_contextual: bool = True):
|
||||
"""
|
||||
Initialize contextual agent.
|
||||
|
||||
Args:
|
||||
config: Configuration object
|
||||
kb_tools: Contextual knowledge base tools
|
||||
use_contextual: Whether to use contextual retrieval
|
||||
"""
|
||||
self.use_contextual = use_contextual
|
||||
|
||||
# Initialize base class but skip KB tools initialization
|
||||
self.config = config or Config.from_env()
|
||||
self._init_llm_client()
|
||||
|
||||
# Use provided contextual KB tools or create new ones
|
||||
if kb_tools:
|
||||
self.kb_tools = kb_tools
|
||||
else:
|
||||
self.kb_tools = ContextualKnowledgeBaseTools(
|
||||
self.config.knowledge_base,
|
||||
use_contextual=use_contextual,
|
||||
enable_comparison=False
|
||||
)
|
||||
|
||||
# Initialize other components
|
||||
self.conversation_history = []
|
||||
self.tools = self._get_contextual_tool_definitions()
|
||||
|
||||
# Track retrieval metrics
|
||||
self.retrieval_metrics = {
|
||||
"queries_executed": 0,
|
||||
"total_chunks_retrieved": 0,
|
||||
"avg_retrieval_score": 0.0,
|
||||
"contextual_chunks_used": 0,
|
||||
"non_contextual_chunks_used": 0
|
||||
}
|
||||
|
||||
mode = "CONTEXTUAL" if use_contextual else "NON-CONTEXTUAL"
|
||||
logger.info(f"Initialized ContextualAgenticRAG in {mode} mode")
|
||||
|
||||
def _get_system_prompt(self) -> str:
|
||||
"""Enhanced system prompt that leverages contextual information"""
|
||||
base_prompt = super()._get_system_prompt()
|
||||
|
||||
if self.use_contextual:
|
||||
contextual_addition = """
|
||||
|
||||
## Using Contextual Information:
|
||||
|
||||
When you receive search results, pay attention to the contextual information provided. Each chunk may include:
|
||||
- Context that situates the chunk within the larger document
|
||||
- Information about what section or topic the chunk belongs to
|
||||
- Related entities and concepts mentioned
|
||||
|
||||
Use this contextual information to:
|
||||
1. Better understand the relevance of each chunk
|
||||
2. Identify connections between different pieces of information
|
||||
3. Provide more accurate and complete answers
|
||||
|
||||
Remember that contextual chunks have been enhanced with additional information to improve retrieval accuracy."""
|
||||
|
||||
return base_prompt + contextual_addition
|
||||
|
||||
return base_prompt
|
||||
|
||||
def _execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
|
||||
"""
|
||||
Execute tool with enhanced logging for educational purposes.
|
||||
|
||||
This override adds detailed logging to show how contextual
|
||||
retrieval improves the search results.
|
||||
"""
|
||||
logger.debug(f"Executing tool: {tool_name} with args: {arguments}")
|
||||
|
||||
try:
|
||||
if tool_name == "contextual_knowledge_search":
|
||||
query = arguments.get("query", "")
|
||||
method = arguments.get("method", "hybrid")
|
||||
top_k = arguments.get("top_k", 20)
|
||||
|
||||
logger.info(f"\n{'='*60}")
|
||||
logger.info(f"CONTEXTUAL SEARCH EXECUTION")
|
||||
logger.info(f"{'='*60}")
|
||||
logger.info(f"Query: {query}")
|
||||
logger.info(f"Method: {method}")
|
||||
logger.info(f"Top-K: {top_k}")
|
||||
logger.info(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
|
||||
|
||||
# Perform search
|
||||
results = self.kb_tools.contextual_search(query, method, top_k)
|
||||
|
||||
if not results:
|
||||
logger.info("No results found")
|
||||
return {"status": "no_results", "message": "No relevant documents found"}
|
||||
|
||||
# Log retrieval quality
|
||||
avg_score = sum(r.score for r in results) / len(results)
|
||||
logger.info(f"\nRetrieval Statistics:")
|
||||
logger.info(f" Results found: {len(results)}")
|
||||
logger.info(f" Average score: {avg_score:.4f}")
|
||||
logger.info(f" Top score: {results[0].score:.4f}")
|
||||
|
||||
# Show distribution of retrieval methods if hybrid
|
||||
if method == "hybrid":
|
||||
bm25_only = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score == 0)
|
||||
embedding_only = sum(1 for r in results if r.embedding_score > 0 and r.bm25_score == 0)
|
||||
both = sum(1 for r in results if r.bm25_score > 0 and r.embedding_score > 0)
|
||||
|
||||
logger.info(f"\nRetrieval Method Distribution:")
|
||||
logger.info(f" BM25 only: {bm25_only}")
|
||||
logger.info(f" Embedding only: {embedding_only}")
|
||||
logger.info(f" Both methods: {both}")
|
||||
|
||||
# Update metrics
|
||||
self.retrieval_metrics["queries_executed"] += 1
|
||||
self.retrieval_metrics["total_chunks_retrieved"] += len(results)
|
||||
|
||||
# Calculate running average score
|
||||
prev_avg = self.retrieval_metrics["avg_retrieval_score"]
|
||||
n = self.retrieval_metrics["queries_executed"]
|
||||
self.retrieval_metrics["avg_retrieval_score"] = (
|
||||
(prev_avg * (n - 1) + avg_score) / n
|
||||
)
|
||||
|
||||
if self.use_contextual:
|
||||
self.retrieval_metrics["contextual_chunks_used"] += len(results)
|
||||
else:
|
||||
self.retrieval_metrics["non_contextual_chunks_used"] += len(results)
|
||||
|
||||
# Format results for agent
|
||||
formatted_results = []
|
||||
for i, r in enumerate(results, 1):
|
||||
result_dict = {
|
||||
"rank": i,
|
||||
"doc_id": r.doc_id,
|
||||
"chunk_id": r.chunk_id,
|
||||
"text": r.text,
|
||||
"score": r.score
|
||||
}
|
||||
|
||||
formatted_results.append(result_dict)
|
||||
|
||||
# Log top results
|
||||
if i <= 3:
|
||||
logger.info(f"\nTop Result {i}:")
|
||||
logger.info(f" Score: {r.score:.4f}")
|
||||
logger.info(f" Text: {r.text}")
|
||||
|
||||
logger.info(f"{'='*60}\n")
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
"results": formatted_results,
|
||||
"total_found": len(results),
|
||||
"avg_score": avg_score,
|
||||
"search_method": method,
|
||||
"is_contextual": self.use_contextual
|
||||
}
|
||||
|
||||
elif tool_name == "get_document":
|
||||
# Use parent implementation
|
||||
return super()._execute_tool(tool_name, arguments)
|
||||
|
||||
else:
|
||||
return {"status": "error", "message": f"Unknown tool: {tool_name}"}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Tool execution error: {e}")
|
||||
return {"status": "error", "message": str(e)}
|
||||
|
||||
def _get_contextual_tool_definitions(self) -> List[Dict[str, Any]]:
|
||||
"""Get tool definitions with contextual search capabilities"""
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "contextual_knowledge_search",
|
||||
"description": "Search the knowledge base using contextual or non-contextual retrieval. Returns chunks with optional context information that helps understand their relevance.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"query": {
|
||||
"type": "string",
|
||||
"description": "Natural language search query"
|
||||
},
|
||||
"method": {
|
||||
"type": "string",
|
||||
"enum": ["bm25", "embedding", "hybrid"],
|
||||
"description": "Search method to use (default: hybrid)",
|
||||
"default": "hybrid"
|
||||
},
|
||||
"top_k": {
|
||||
"type": "integer",
|
||||
"description": "Number of results to retrieve (default: 20)",
|
||||
"default": 20
|
||||
}
|
||||
},
|
||||
"required": ["query"]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_document",
|
||||
"description": "Retrieve the complete content of a specific document",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"doc_id": {
|
||||
"type": "string",
|
||||
"description": "Document ID to retrieve"
|
||||
}
|
||||
},
|
||||
"required": ["doc_id"]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
def query(self, user_query: str, stream: bool = None) -> Any:
|
||||
"""
|
||||
Process query with contextual retrieval and detailed logging.
|
||||
|
||||
This override adds educational logging to show the retrieval process.
|
||||
"""
|
||||
logger.info(f"\n{'='*80}")
|
||||
logger.info(f"PROCESSING QUERY: {user_query}")
|
||||
logger.info(f"Mode: {'CONTEXTUAL' if self.use_contextual else 'NON-CONTEXTUAL'}")
|
||||
logger.info(f"{'='*80}\n")
|
||||
|
||||
# Call parent implementation
|
||||
result = super().query(user_query, stream)
|
||||
|
||||
# Log final metrics
|
||||
logger.info(f"\n{'='*60}")
|
||||
logger.info("QUERY PROCESSING COMPLETE")
|
||||
logger.info(f"{'='*60}")
|
||||
logger.info("Retrieval Metrics:")
|
||||
logger.info(f" Queries executed: {self.retrieval_metrics['queries_executed']}")
|
||||
logger.info(f" Total chunks retrieved: {self.retrieval_metrics['total_chunks_retrieved']}")
|
||||
logger.info(f" Average retrieval score: {self.retrieval_metrics['avg_retrieval_score']:.4f}")
|
||||
|
||||
if self.use_contextual:
|
||||
logger.info(f" Contextual chunks used: {self.retrieval_metrics['contextual_chunks_used']}")
|
||||
else:
|
||||
logger.info(f" Non-contextual chunks used: {self.retrieval_metrics['non_contextual_chunks_used']}")
|
||||
|
||||
logger.info(f"{'='*60}\n")
|
||||
|
||||
return result
|
||||
|
||||
def get_retrieval_metrics(self) -> Dict[str, Any]:
|
||||
"""Get detailed retrieval metrics for analysis"""
|
||||
metrics = self.retrieval_metrics.copy()
|
||||
|
||||
# Add KB statistics
|
||||
kb_stats = self.kb_tools.get_statistics()
|
||||
metrics["kb_stats"] = kb_stats
|
||||
|
||||
# Add mode information
|
||||
metrics["mode"] = "contextual" if self.use_contextual else "non_contextual"
|
||||
metrics["timestamp"] = datetime.now().isoformat()
|
||||
|
||||
return metrics
|
||||
|
||||
def compare_modes(self, user_query: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Compare contextual vs non-contextual modes for the same query.
|
||||
|
||||
Educational method to demonstrate the difference in retrieval quality.
|
||||
"""
|
||||
comparison = {
|
||||
"query": user_query,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"modes": {}
|
||||
}
|
||||
|
||||
# Test contextual mode
|
||||
logger.info("Testing CONTEXTUAL mode...")
|
||||
self.use_contextual = True
|
||||
self.kb_tools.use_contextual = True
|
||||
contextual_response = self.query(user_query, stream=False)
|
||||
comparison["modes"]["contextual"] = {
|
||||
"response": contextual_response,
|
||||
"metrics": self.get_retrieval_metrics()
|
||||
}
|
||||
|
||||
# Reset metrics
|
||||
self.retrieval_metrics = {
|
||||
"queries_executed": 0,
|
||||
"total_chunks_retrieved": 0,
|
||||
"avg_retrieval_score": 0.0,
|
||||
"contextual_chunks_used": 0,
|
||||
"non_contextual_chunks_used": 0
|
||||
}
|
||||
|
||||
# Test non-contextual mode
|
||||
logger.info("Testing NON-CONTEXTUAL mode...")
|
||||
self.use_contextual = False
|
||||
self.kb_tools.use_contextual = False
|
||||
non_contextual_response = self.query(user_query, stream=False)
|
||||
comparison["modes"]["non_contextual"] = {
|
||||
"response": non_contextual_response,
|
||||
"metrics": self.get_retrieval_metrics()
|
||||
}
|
||||
|
||||
# Analyze differences
|
||||
contextual_score = comparison["modes"]["contextual"]["metrics"]["avg_retrieval_score"]
|
||||
non_contextual_score = comparison["modes"]["non_contextual"]["metrics"]["avg_retrieval_score"]
|
||||
|
||||
improvement = ((contextual_score - non_contextual_score) / non_contextual_score * 100) if non_contextual_score > 0 else 0
|
||||
|
||||
comparison["analysis"] = {
|
||||
"score_improvement_pct": improvement,
|
||||
"contextual_avg_score": contextual_score,
|
||||
"non_contextual_avg_score": non_contextual_score,
|
||||
"recommendation": (
|
||||
"Use contextual retrieval" if improvement > 10
|
||||
else "Consider cost-benefit of contextual retrieval"
|
||||
)
|
||||
}
|
||||
|
||||
return comparison
|
||||
@@ -0,0 +1,536 @@
|
||||
"""Contextual Chunking Module - Educational implementation of Anthropic's Contextual Retrieval
|
||||
|
||||
This module demonstrates the key insight from Anthropic's research:
|
||||
- Traditional RAG destroys context by chunking documents
|
||||
- Contextual Retrieval prepends chunk-specific context before embedding
|
||||
- This preserves semantic meaning that would otherwise be lost
|
||||
"""
|
||||
|
||||
import json
|
||||
import hashlib
|
||||
import logging
|
||||
import requests
|
||||
from typing import List, Dict, Any, Optional, Tuple
|
||||
from pathlib import Path
|
||||
from datetime import datetime
|
||||
from dataclasses import dataclass, field
|
||||
import time
|
||||
from openai import OpenAI
|
||||
from config import ChunkingConfig, KnowledgeBaseConfig, KnowledgeBaseType, LLMConfig
|
||||
|
||||
|
||||
def _reasoning_safe_temperature(model, requested=1.0):
|
||||
"""Reasoning models (Kimi K3, GPT-5, ...) only accept temperature=1.
|
||||
Return 1 for those; otherwise the requested value so non-reasoning
|
||||
providers (Doubao, DeepSeek, older Moonshot) are unchanged."""
|
||||
m = str(model or "").lower().replace("/", "-")
|
||||
return 1 if ("kimi-k3" in m or "gpt-5" in m) else requested
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextualChunk:
|
||||
"""Enhanced chunk with contextual information"""
|
||||
chunk_id: str
|
||||
doc_id: str
|
||||
text: str # Original chunk text
|
||||
context: str # Generated contextual description
|
||||
contextualized_text: str # Context + original text
|
||||
chunk_index: int
|
||||
char_count: int
|
||||
context_tokens: int = 0 # Track token usage
|
||||
generation_time: float = 0.0 # Track generation time
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"chunk_id": self.chunk_id,
|
||||
"doc_id": self.doc_id,
|
||||
"text": self.text,
|
||||
"context": self.context,
|
||||
"contextualized_text": self.contextualized_text,
|
||||
"chunk_index": self.chunk_index,
|
||||
"char_count": self.char_count,
|
||||
"context_tokens": self.context_tokens,
|
||||
"generation_time": self.generation_time,
|
||||
"metadata": self.metadata
|
||||
}
|
||||
|
||||
|
||||
class ContextualChunker:
|
||||
"""
|
||||
Implements contextual chunking inspired by Anthropic's Contextual Retrieval.
|
||||
|
||||
Key Educational Points:
|
||||
1. Context Generation: Uses LLM to generate chunk-specific context
|
||||
2. Prepending Strategy: Context is prepended to chunks before embedding
|
||||
3. BM25 Enhancement: Contextual chunks improve both semantic and lexical search
|
||||
4. Cost Optimization: Uses caching strategies to reduce API costs
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
chunking_config: Optional[ChunkingConfig] = None,
|
||||
llm_config: Optional[LLMConfig] = None,
|
||||
use_contextual: bool = True):
|
||||
"""
|
||||
Initialize the contextual chunker.
|
||||
|
||||
Args:
|
||||
chunking_config: Configuration for chunking parameters
|
||||
llm_config: LLM configuration for context generation
|
||||
use_contextual: Whether to generate contextual chunks (for comparison)
|
||||
"""
|
||||
self.chunking_config = chunking_config or ChunkingConfig()
|
||||
self.llm_config = llm_config or LLMConfig()
|
||||
self.use_contextual = use_contextual
|
||||
|
||||
# Initialize LLM client for context generation
|
||||
if self.use_contextual:
|
||||
self._init_llm_client()
|
||||
|
||||
# Statistics tracking
|
||||
self.stats = {
|
||||
"total_chunks": 0,
|
||||
"contextual_chunks": 0,
|
||||
"total_context_tokens": 0,
|
||||
"total_generation_time": 0.0,
|
||||
"cache_hits": 0,
|
||||
"cache_misses": 0
|
||||
}
|
||||
|
||||
# Context cache to avoid regenerating for similar chunks
|
||||
self.context_cache = {}
|
||||
|
||||
logger.info(f"Initialized ContextualChunker (contextual={use_contextual})")
|
||||
|
||||
def _init_llm_client(self):
|
||||
"""Initialize LLM client for context generation"""
|
||||
client_config, model = self.llm_config.get_client_config()
|
||||
base_url = client_config.pop("base_url", None)
|
||||
|
||||
if base_url:
|
||||
self.client = OpenAI(base_url=base_url, **client_config)
|
||||
else:
|
||||
self.client = OpenAI(**client_config)
|
||||
|
||||
self.model = model
|
||||
logger.info(f"Using {self.llm_config.provider} ({self.model}) for context generation")
|
||||
|
||||
def chunk_document(self,
|
||||
text: str,
|
||||
doc_id: str,
|
||||
doc_metadata: Optional[Dict[str, Any]] = None,
|
||||
on_chunk_ready: Optional[callable] = None) -> List[ContextualChunk]:
|
||||
"""
|
||||
Chunk a document with optional contextual enhancement.
|
||||
|
||||
Educational Note:
|
||||
This is the core innovation - each chunk gets contextualized
|
||||
with information about its position and meaning within the document.
|
||||
|
||||
Args:
|
||||
text: Full document text
|
||||
doc_id: Document identifier
|
||||
doc_metadata: Optional document metadata
|
||||
|
||||
Returns:
|
||||
List of ContextualChunk objects
|
||||
"""
|
||||
logger.info(f"Starting chunking for document {doc_id}")
|
||||
start_time = time.time()
|
||||
|
||||
# Step 1: Create basic chunks
|
||||
basic_chunks = self._create_basic_chunks(text, doc_id)
|
||||
logger.info(f"Created {len(basic_chunks)} basic chunks")
|
||||
|
||||
# Step 2: Generate contextual enhancements if enabled
|
||||
if self.use_contextual:
|
||||
contextual_chunks = self._generate_contextual_chunks(
|
||||
basic_chunks, text, doc_id, doc_metadata, on_chunk_ready
|
||||
)
|
||||
else:
|
||||
# Create non-contextual chunks for comparison
|
||||
contextual_chunks = []
|
||||
for chunk in basic_chunks:
|
||||
contextual_chunks.append(ContextualChunk(
|
||||
chunk_id=chunk["chunk_id"],
|
||||
doc_id=chunk["doc_id"],
|
||||
text=chunk["text"],
|
||||
context="", # No context in non-contextual mode
|
||||
contextualized_text=chunk["text"], # Just the original text
|
||||
chunk_index=chunk["chunk_index"],
|
||||
char_count=chunk["char_count"],
|
||||
metadata={"contextual": False}
|
||||
))
|
||||
|
||||
# Update statistics
|
||||
self.stats["total_chunks"] += len(contextual_chunks)
|
||||
if self.use_contextual:
|
||||
self.stats["contextual_chunks"] += len(contextual_chunks)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
logger.info(f"Chunking completed in {elapsed:.2f}s")
|
||||
logger.info(f"Statistics: {json.dumps(self.stats, indent=2)}")
|
||||
|
||||
return contextual_chunks
|
||||
|
||||
def _create_basic_chunks(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
|
||||
"""Create basic chunks using traditional chunking"""
|
||||
chunks = []
|
||||
|
||||
if self.chunking_config.respect_paragraph_boundary:
|
||||
chunks = self._chunk_by_paragraphs(text, doc_id)
|
||||
else:
|
||||
chunks = self._chunk_by_size(text, doc_id)
|
||||
|
||||
return chunks
|
||||
|
||||
def _chunk_by_paragraphs(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
|
||||
"""Chunk text respecting paragraph boundaries"""
|
||||
paragraphs = text.split('\n\n')
|
||||
chunks = []
|
||||
current_chunk = []
|
||||
current_size = 0
|
||||
|
||||
for para in paragraphs:
|
||||
para = para.strip()
|
||||
if not para:
|
||||
continue
|
||||
|
||||
para_size = len(para)
|
||||
|
||||
# Handle oversized paragraphs
|
||||
if para_size > self.chunking_config.max_chunk_size:
|
||||
if current_chunk:
|
||||
chunk_text = '\n\n'.join(current_chunk)
|
||||
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
|
||||
current_chunk = []
|
||||
current_size = 0
|
||||
|
||||
# Split large paragraph into sentences
|
||||
sentences = self._split_into_sentences(para)
|
||||
for sent in sentences:
|
||||
if len(sent) > self.chunking_config.max_chunk_size:
|
||||
# Force split very long sentences
|
||||
for i in range(0, len(sent), self.chunking_config.chunk_size):
|
||||
sub_chunk = sent[i:i + self.chunking_config.chunk_size]
|
||||
chunks.append(self._create_basic_chunk(sub_chunk, doc_id, len(chunks)))
|
||||
else:
|
||||
chunks.append(self._create_basic_chunk(sent, doc_id, len(chunks)))
|
||||
continue
|
||||
|
||||
# Check if adding this paragraph exceeds chunk size
|
||||
if current_size + para_size > self.chunking_config.chunk_size and current_chunk:
|
||||
chunk_text = '\n\n'.join(current_chunk)
|
||||
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
|
||||
|
||||
# Start new chunk with overlap
|
||||
if self.chunking_config.chunk_overlap > 0 and current_chunk:
|
||||
current_chunk = [current_chunk[-1], para]
|
||||
current_size = len(current_chunk[0]) + para_size
|
||||
else:
|
||||
current_chunk = [para]
|
||||
current_size = para_size
|
||||
else:
|
||||
current_chunk.append(para)
|
||||
current_size += para_size
|
||||
|
||||
# Save final chunk
|
||||
if current_chunk:
|
||||
chunk_text = '\n\n'.join(current_chunk)
|
||||
if len(chunk_text) >= self.chunking_config.min_chunk_size:
|
||||
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
|
||||
|
||||
return chunks
|
||||
|
||||
def _chunk_by_size(self, text: str, doc_id: str) -> List[Dict[str, Any]]:
|
||||
"""Simple size-based chunking"""
|
||||
chunks = []
|
||||
|
||||
for i in range(0, len(text), self.chunking_config.chunk_size - self.chunking_config.chunk_overlap):
|
||||
chunk_text = text[i:i + self.chunking_config.chunk_size]
|
||||
|
||||
if len(chunk_text) >= self.chunking_config.min_chunk_size:
|
||||
chunks.append(self._create_basic_chunk(chunk_text, doc_id, len(chunks)))
|
||||
|
||||
return chunks
|
||||
|
||||
def _split_into_sentences(self, text: str) -> List[str]:
|
||||
"""Split text into sentences"""
|
||||
import re
|
||||
|
||||
# Handle both English and Chinese sentence endings
|
||||
sentences = re.split(r'([。!?\.!?]+)', text)
|
||||
|
||||
# Reconstruct sentences with their endings
|
||||
result = []
|
||||
# Step to the end of the list: re.split with a capturing group yields
|
||||
# [text, delim, text, delim, ..., trailing_text], so stopping at
|
||||
# len(sentences) - 1 dropped the trailing fragment whenever the text
|
||||
# did not end in terminal punctuation (and returned [] for text with
|
||||
# none at all). The strip-and-filter below still discards the empty
|
||||
# tail that re.split produces when the text does end in punctuation.
|
||||
for i in range(0, len(sentences), 2):
|
||||
if i + 1 < len(sentences):
|
||||
result.append(sentences[i] + sentences[i + 1])
|
||||
else:
|
||||
result.append(sentences[i])
|
||||
|
||||
return [s.strip() for s in result if s.strip()]
|
||||
|
||||
def _create_basic_chunk(self, text: str, doc_id: str, chunk_index: int) -> Dict[str, Any]:
|
||||
"""Create a basic chunk dictionary"""
|
||||
chunk_id = f"{doc_id}_chunk_{chunk_index}"
|
||||
|
||||
return {
|
||||
"chunk_id": chunk_id,
|
||||
"doc_id": doc_id,
|
||||
"text": text,
|
||||
"chunk_index": chunk_index,
|
||||
"char_count": len(text),
|
||||
"hash": hashlib.md5(text.encode()).hexdigest()
|
||||
}
|
||||
|
||||
def _generate_contextual_chunks(self,
|
||||
basic_chunks: List[Dict[str, Any]],
|
||||
full_document: str,
|
||||
doc_id: str,
|
||||
doc_metadata: Optional[Dict[str, Any]] = None,
|
||||
on_chunk_ready: Optional[callable] = None) -> List[ContextualChunk]:
|
||||
"""
|
||||
Generate contextual enhancements for chunks.
|
||||
|
||||
Educational Note:
|
||||
Following Anthropic's Contextual Retrieval approach:
|
||||
- Each chunk gets a concise context explaining its position in the document
|
||||
- Context is prepended to the chunk before embedding
|
||||
- This dramatically improves retrieval accuracy
|
||||
"""
|
||||
contextual_chunks = []
|
||||
|
||||
# No document summary needed - Anthropic's approach doesn't use it
|
||||
doc_summary = None
|
||||
|
||||
for i, chunk in enumerate(basic_chunks):
|
||||
logger.info(f"Generating context for chunk {i+1}/{len(basic_chunks)}")
|
||||
|
||||
# Check cache first
|
||||
chunk_hash = chunk["hash"]
|
||||
if chunk_hash in self.context_cache:
|
||||
context = self.context_cache[chunk_hash]
|
||||
generation_time = 0.0
|
||||
context_tokens = 0
|
||||
self.stats["cache_hits"] += 1
|
||||
logger.debug(f"Cache hit for chunk {chunk['chunk_id']}")
|
||||
else:
|
||||
# Generate new context using Anthropic's approach
|
||||
context, context_tokens, generation_time = self._generate_chunk_context(
|
||||
chunk["text"],
|
||||
full_document
|
||||
)
|
||||
|
||||
# Cache the context
|
||||
self.context_cache[chunk_hash] = context
|
||||
self.stats["cache_misses"] += 1
|
||||
self.stats["total_context_tokens"] += context_tokens
|
||||
self.stats["total_generation_time"] += generation_time
|
||||
|
||||
# Create contextual chunk
|
||||
contextualized_text = f"{context}\n\n{chunk['text']}" if context else chunk["text"]
|
||||
|
||||
contextual_chunk = ContextualChunk(
|
||||
chunk_id=chunk["chunk_id"],
|
||||
doc_id=doc_id,
|
||||
text=chunk["text"],
|
||||
context=context,
|
||||
contextualized_text=contextualized_text,
|
||||
chunk_index=chunk["chunk_index"],
|
||||
char_count=len(contextualized_text),
|
||||
context_tokens=context_tokens,
|
||||
generation_time=generation_time,
|
||||
metadata={
|
||||
"contextual": True,
|
||||
"original_char_count": chunk["char_count"],
|
||||
"context_char_count": len(context)
|
||||
}
|
||||
)
|
||||
|
||||
contextual_chunks.append(contextual_chunk)
|
||||
|
||||
# Call the callback immediately if provided
|
||||
if on_chunk_ready:
|
||||
try:
|
||||
on_chunk_ready(contextual_chunk)
|
||||
except Exception as e:
|
||||
logger.error(f"Error in on_chunk_ready callback: {e}")
|
||||
|
||||
# Log progress
|
||||
if (i + 1) % 10 == 0:
|
||||
avg_time = self.stats["total_generation_time"] / (i + 1)
|
||||
logger.info(f"Progress: {i+1}/{len(basic_chunks)} chunks, avg time: {avg_time:.2f}s")
|
||||
|
||||
return contextual_chunks
|
||||
|
||||
def _generate_document_summary(self, document: str, doc_id: str) -> str:
|
||||
"""
|
||||
DEPRECATED: Not used in Anthropic's Contextual Retrieval approach.
|
||||
|
||||
Educational Note:
|
||||
Anthropic's research shows that document summaries don't significantly
|
||||
improve retrieval. Instead, they provide the full document directly
|
||||
when generating chunk-specific context. This allows the LLM to understand
|
||||
the exact context needed for each specific chunk.
|
||||
"""
|
||||
# This method is kept for backward compatibility but returns empty string
|
||||
return ""
|
||||
|
||||
def _generate_chunk_context(self,
|
||||
chunk_text: str,
|
||||
full_document: str) -> Tuple[str, int, float]:
|
||||
"""
|
||||
Generate contextual description for a chunk using Anthropic's exact approach.
|
||||
|
||||
This follows Anthropic's Contextual Retrieval template exactly:
|
||||
1. Provide the full document
|
||||
2. Show the specific chunk
|
||||
3. Ask for concise context to situate the chunk
|
||||
|
||||
Returns:
|
||||
Tuple of (context, token_count, generation_time)
|
||||
"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# Use the exact prompt from Anthropic's blog post
|
||||
# with added instruction to use the same language as the document
|
||||
prompt = f"""<document>
|
||||
{full_document}
|
||||
</document>
|
||||
|
||||
Here is the chunk we want to situate within the whole document
|
||||
|
||||
<chunk>
|
||||
{chunk_text}
|
||||
</chunk>
|
||||
|
||||
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else. You MUST use the same language as the document."""
|
||||
|
||||
# Use the exact approach from Anthropic - no system message needed
|
||||
response = self.client.chat.completions.create(
|
||||
model=self.model,
|
||||
messages=[
|
||||
{"role": "user", "content": prompt}
|
||||
],
|
||||
temperature=_reasoning_safe_temperature(self.model, 0.3), # Low temperature for consistency
|
||||
max_tokens=100 # Anthropic mentions 50-100 tokens typically
|
||||
)
|
||||
|
||||
context = response.choices[0].message.content.strip()
|
||||
|
||||
# Estimate token count (rough approximation)
|
||||
token_count = len(prompt.split()) + len(context.split())
|
||||
generation_time = time.time() - start_time
|
||||
|
||||
logger.info(f"Generated context in {generation_time:.2f}s: {context}")
|
||||
|
||||
return context, token_count, generation_time
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error generating chunk context: {e}")
|
||||
return "", 0, time.time() - start_time
|
||||
|
||||
def compare_retrieval_methods(self,
|
||||
query: str,
|
||||
contextual_chunks: List[ContextualChunk],
|
||||
non_contextual_chunks: List[ContextualChunk],
|
||||
top_k: int = 5) -> Dict[str, Any]:
|
||||
"""
|
||||
Compare contextual vs non-contextual retrieval on the SAME query.
|
||||
|
||||
Educational Note:
|
||||
This is the ``compare_retrieval_methods`` capability referenced in
|
||||
实验 3-10. It builds two BM25 indexes fully offline (no API / server):
|
||||
* contextual index over ``contextualized_text`` (前缀 + 原文)
|
||||
* plain index over the original chunk ``text``
|
||||
and returns the top-k ranked chunks under each, so the caller can see
|
||||
exactly how the contextual prefix re-ranks the same corpus.
|
||||
"""
|
||||
from rank_bm25 import BM25Okapi
|
||||
import numpy as np
|
||||
from compare_retrieval import tokenize
|
||||
|
||||
results = {
|
||||
"query": query,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"contextual_results": [],
|
||||
"non_contextual_results": [],
|
||||
"analysis": {}
|
||||
}
|
||||
|
||||
def _rank(chunks: List[ContextualChunk], field: str):
|
||||
if not chunks:
|
||||
return []
|
||||
corpus = [tokenize(getattr(c, field)) for c in chunks]
|
||||
index = BM25Okapi(corpus)
|
||||
scores = index.get_scores(tokenize(query))
|
||||
order = np.argsort(scores)[::-1][:top_k]
|
||||
ranked = []
|
||||
for rank, idx in enumerate(order, 1):
|
||||
c = chunks[idx]
|
||||
ranked.append({
|
||||
"chunk_id": c.chunk_id,
|
||||
"score": float(scores[idx]),
|
||||
"rank": rank,
|
||||
"text": c.text[:200],
|
||||
"context": c.context[:200],
|
||||
})
|
||||
return ranked
|
||||
|
||||
results["contextual_results"] = _rank(contextual_chunks, "contextualized_text")
|
||||
results["non_contextual_results"] = _rank(non_contextual_chunks, "text")
|
||||
|
||||
ctx_top = results["contextual_results"][0] if results["contextual_results"] else None
|
||||
plain_top = results["non_contextual_results"][0] if results["non_contextual_results"] else None
|
||||
results["analysis"] = {
|
||||
"contextual_top_chunk": ctx_top["chunk_id"] if ctx_top else None,
|
||||
"non_contextual_top_chunk": plain_top["chunk_id"] if plain_top else None,
|
||||
"contextual_top_score": ctx_top["score"] if ctx_top else 0.0,
|
||||
"non_contextual_top_score": plain_top["score"] if plain_top else 0.0,
|
||||
"top1_changed": bool(ctx_top and plain_top and ctx_top["chunk_id"] != plain_top["chunk_id"]),
|
||||
}
|
||||
|
||||
logger.info(f"Compared retrieval methods for query: {query} | "
|
||||
f"top1 changed={results['analysis']['top1_changed']}")
|
||||
|
||||
return results
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get chunking statistics"""
|
||||
stats = self.stats.copy()
|
||||
|
||||
# Calculate averages
|
||||
if stats["contextual_chunks"] > 0:
|
||||
stats["avg_context_tokens"] = stats["total_context_tokens"] / stats["contextual_chunks"]
|
||||
stats["avg_generation_time"] = stats["total_generation_time"] / stats["contextual_chunks"]
|
||||
else:
|
||||
stats["avg_context_tokens"] = 0
|
||||
stats["avg_generation_time"] = 0
|
||||
|
||||
# Cache efficiency
|
||||
total_cache_ops = stats["cache_hits"] + stats["cache_misses"]
|
||||
if total_cache_ops > 0:
|
||||
stats["cache_hit_rate"] = stats["cache_hits"] / total_cache_ops
|
||||
else:
|
||||
stats["cache_hit_rate"] = 0
|
||||
|
||||
# Cost estimation
|
||||
if self.llm_config.provider == "openai":
|
||||
cost_per_1k = 0.03
|
||||
else:
|
||||
cost_per_1k = 0.01
|
||||
|
||||
stats["estimated_cost"] = (stats["total_context_tokens"] / 1000) * cost_per_1k
|
||||
|
||||
return stats
|
||||
@@ -0,0 +1,658 @@
|
||||
"""Enhanced tools for contextual retrieval with BM25 and semantic search
|
||||
|
||||
Educational implementation showing how contextual chunks improve both
|
||||
BM25 (lexical) and embedding (semantic) retrieval.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import requests
|
||||
import numpy as np
|
||||
from typing import Dict, Any, List, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
import time
|
||||
from rank_bm25 import BM25Okapi
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
|
||||
from config import KnowledgeBaseConfig, KnowledgeBaseType
|
||||
from tools import KnowledgeBaseTools, SearchResult
|
||||
from contextual_chunking import ContextualChunk
|
||||
# Shared tokenizer: 中文没有空格,原先的 .lower().split() 会把整段当成一个 token,
|
||||
# 导致 BM25 在中文语料上几乎失效。统一改用 compare_retrieval.tokenize(jieba 分词)。
|
||||
from compare_retrieval import tokenize as _bm25_tokenize
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextualSearchResult(SearchResult):
|
||||
"""Enhanced search result with contextual information"""
|
||||
is_contextual: bool = False
|
||||
context_text: str = ""
|
||||
bm25_score: float = 0.0
|
||||
embedding_score: float = 0.0
|
||||
hybrid_score: float = 0.0
|
||||
retrieval_method: str = "hybrid" # bm25, embedding, or hybrid
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
base = super().to_dict()
|
||||
base.update({
|
||||
"is_contextual": self.is_contextual,
|
||||
"context_text": self.context_text,
|
||||
"bm25_score": self.bm25_score,
|
||||
"embedding_score": self.embedding_score,
|
||||
"hybrid_score": self.hybrid_score,
|
||||
"retrieval_method": self.retrieval_method
|
||||
})
|
||||
return base
|
||||
|
||||
|
||||
class ContextualKnowledgeBaseTools(KnowledgeBaseTools):
|
||||
"""
|
||||
Enhanced knowledge base tools with contextual retrieval support.
|
||||
|
||||
Key Educational Points:
|
||||
1. Dual Indexing: Maintains both contextual and non-contextual indexes
|
||||
2. BM25 Enhancement: Shows how context improves lexical matching
|
||||
3. Hybrid Search: Combines BM25 and semantic search with rank fusion
|
||||
4. Comparison Mode: Allows side-by-side evaluation of methods
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
config: KnowledgeBaseConfig,
|
||||
use_contextual: bool = True,
|
||||
enable_comparison: bool = False):
|
||||
"""
|
||||
Initialize contextual knowledge base tools.
|
||||
|
||||
Args:
|
||||
config: Knowledge base configuration
|
||||
use_contextual: Whether to use contextual retrieval
|
||||
enable_comparison: Whether to enable comparison mode
|
||||
"""
|
||||
super().__init__(config)
|
||||
self.use_contextual = use_contextual
|
||||
self.enable_comparison = enable_comparison
|
||||
|
||||
# BM25 indexes for lexical search
|
||||
self.bm25_index = None
|
||||
self.bm25_contextual_index = None
|
||||
self.bm25_corpus = []
|
||||
self.bm25_contextual_corpus = []
|
||||
|
||||
# Document and chunk storage
|
||||
self.chunk_store = {} # chunk_id -> ContextualChunk
|
||||
self.contextual_chunk_store = {} # chunk_id -> ContextualChunk (with context)
|
||||
|
||||
# Index paths
|
||||
self.index_dir = Path("indexes")
|
||||
self.index_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Load existing indexes if available
|
||||
self._load_indexes()
|
||||
|
||||
# Statistics
|
||||
self.search_stats = {
|
||||
"total_searches": 0,
|
||||
"contextual_searches": 0,
|
||||
"non_contextual_searches": 0,
|
||||
"comparison_searches": 0,
|
||||
"avg_retrieval_time": 0.0,
|
||||
"total_retrieval_time": 0.0
|
||||
}
|
||||
|
||||
logger.info(f"Initialized ContextualKnowledgeBaseTools (contextual={use_contextual}, comparison={enable_comparison})")
|
||||
|
||||
def index_contextual_chunks(self, chunks: List[ContextualChunk], rebuild_bm25: bool = True):
|
||||
"""
|
||||
Index contextual chunks for both BM25 and semantic search.
|
||||
|
||||
Educational Note:
|
||||
This demonstrates the dual indexing strategy:
|
||||
- BM25 index on contextualized text for better lexical matching
|
||||
- Semantic embeddings on contextualized text for richer meaning
|
||||
"""
|
||||
logger.info(f"Indexing {len(chunks)} contextual chunks")
|
||||
start_time = time.time()
|
||||
|
||||
# Store chunks
|
||||
for chunk in chunks:
|
||||
self.contextual_chunk_store[chunk.chunk_id] = chunk
|
||||
|
||||
# Also store non-contextual version for comparison
|
||||
non_contextual_chunk = ContextualChunk(
|
||||
chunk_id=chunk.chunk_id + "_nc",
|
||||
doc_id=chunk.doc_id,
|
||||
text=chunk.text,
|
||||
context="",
|
||||
contextualized_text=chunk.text,
|
||||
chunk_index=chunk.chunk_index,
|
||||
char_count=len(chunk.text),
|
||||
metadata={"contextual": False}
|
||||
)
|
||||
self.chunk_store[non_contextual_chunk.chunk_id] = non_contextual_chunk
|
||||
|
||||
# Build BM25 indexes
|
||||
if rebuild_bm25:
|
||||
self._build_bm25_indexes()
|
||||
|
||||
# Index to retrieval pipeline (if local)
|
||||
if self.config.type == KnowledgeBaseType.LOCAL:
|
||||
self._index_chunks_to_pipeline(chunks)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
logger.info(f"Indexed {len(chunks)} chunks in {elapsed:.2f}s")
|
||||
|
||||
# Save indexes
|
||||
self._save_indexes()
|
||||
|
||||
def _build_bm25_indexes(self):
|
||||
"""
|
||||
Build BM25 indexes for both contextual and non-contextual chunks.
|
||||
|
||||
Educational Note:
|
||||
BM25 uses TF-IDF with optimizations:
|
||||
- Term frequency saturation prevents common words from dominating
|
||||
- Document length normalization accounts for varying chunk sizes
|
||||
- The contextual version has richer vocabulary from added context
|
||||
"""
|
||||
logger.info("Building BM25 indexes")
|
||||
|
||||
# Build contextual BM25 index
|
||||
if self.contextual_chunk_store:
|
||||
contextual_texts = []
|
||||
for chunk in self.contextual_chunk_store.values():
|
||||
# Tokenize for BM25 (jieba 中文分词,兼容英文)
|
||||
tokens = _bm25_tokenize(chunk.contextualized_text)
|
||||
contextual_texts.append(tokens)
|
||||
|
||||
self.bm25_contextual_corpus = contextual_texts
|
||||
self.bm25_contextual_index = BM25Okapi(contextual_texts)
|
||||
logger.info(f"Built contextual BM25 index with {len(contextual_texts)} documents")
|
||||
|
||||
# Build non-contextual BM25 index
|
||||
if self.chunk_store:
|
||||
non_contextual_texts = []
|
||||
for chunk in self.chunk_store.values():
|
||||
tokens = _bm25_tokenize(chunk.text)
|
||||
non_contextual_texts.append(tokens)
|
||||
|
||||
self.bm25_corpus = non_contextual_texts
|
||||
self.bm25_index = BM25Okapi(non_contextual_texts)
|
||||
logger.info(f"Built non-contextual BM25 index with {len(non_contextual_texts)} documents")
|
||||
|
||||
def _index_chunks_to_pipeline(self, chunks: List[ContextualChunk]):
|
||||
"""Index chunks to the retrieval pipeline"""
|
||||
for chunk in chunks:
|
||||
try:
|
||||
# Index contextual version
|
||||
if self.use_contextual:
|
||||
response = requests.post(
|
||||
f"{self.config.local_base_url}/index",
|
||||
json={
|
||||
"text": chunk.contextualized_text,
|
||||
"doc_id": chunk.doc_id,
|
||||
"metadata": {
|
||||
"chunk_id": chunk.chunk_id,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"is_contextual": True,
|
||||
"context": chunk.context[:200], # Store truncated context
|
||||
"original_text": chunk.text[:500] # Store truncated original
|
||||
}
|
||||
}, timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
# Also index non-contextual version if in comparison mode
|
||||
if self.enable_comparison:
|
||||
response = requests.post(
|
||||
f"{self.config.local_base_url}/index",
|
||||
json={
|
||||
"text": chunk.text,
|
||||
"doc_id": chunk.doc_id,
|
||||
"metadata": {
|
||||
"chunk_id": chunk.chunk_id + "_nc",
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"is_contextual": False
|
||||
}
|
||||
}, timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error indexing chunk {chunk.chunk_id}: {e}")
|
||||
|
||||
def contextual_search(self,
|
||||
query: str,
|
||||
method: str = "hybrid",
|
||||
top_k: int = 20) -> List[ContextualSearchResult]:
|
||||
"""
|
||||
Perform contextual search using specified method.
|
||||
|
||||
Args:
|
||||
query: Search query
|
||||
method: Search method - "bm25", "embedding", or "hybrid"
|
||||
top_k: Number of results to return
|
||||
|
||||
Educational Note:
|
||||
This demonstrates three retrieval strategies:
|
||||
1. BM25: Pure lexical matching based on term frequency
|
||||
2. Embedding: Semantic similarity using vector embeddings
|
||||
3. Hybrid: Rank fusion combining both approaches
|
||||
"""
|
||||
logger.info(f"Performing {method} search for: {query[:100]}...")
|
||||
start_time = time.time()
|
||||
|
||||
results = []
|
||||
|
||||
if method in ["bm25", "hybrid"]:
|
||||
bm25_results = self._search_bm25(query, self.use_contextual, top_k * 2)
|
||||
results.extend(bm25_results)
|
||||
|
||||
if method in ["embedding", "hybrid"]:
|
||||
embedding_results = self._search_embeddings(query, self.use_contextual, top_k * 2)
|
||||
results.extend(embedding_results)
|
||||
|
||||
if method == "hybrid":
|
||||
# Rank fusion: combine and deduplicate results
|
||||
results = self._rank_fusion(bm25_results, embedding_results, top_k)
|
||||
else:
|
||||
# Sort by score and limit
|
||||
results = sorted(results, key=lambda x: x.score, reverse=True)[:top_k]
|
||||
|
||||
# Update statistics
|
||||
elapsed = time.time() - start_time
|
||||
self.search_stats["total_searches"] += 1
|
||||
if self.use_contextual:
|
||||
self.search_stats["contextual_searches"] += 1
|
||||
else:
|
||||
self.search_stats["non_contextual_searches"] += 1
|
||||
self.search_stats["total_retrieval_time"] += elapsed
|
||||
self.search_stats["avg_retrieval_time"] = (
|
||||
self.search_stats["total_retrieval_time"] / self.search_stats["total_searches"]
|
||||
)
|
||||
|
||||
logger.info(f"Search completed in {elapsed:.2f}s, returned {len(results)} results")
|
||||
|
||||
return results
|
||||
|
||||
def _search_bm25(self, query: str, use_contextual: bool, top_k: int) -> List[ContextualSearchResult]:
|
||||
"""
|
||||
Perform BM25 search.
|
||||
|
||||
Educational Note:
|
||||
BM25 excels at finding exact term matches and handles
|
||||
technical terms, IDs, and specific phrases well.
|
||||
Contextual chunks help by adding synonyms and related terms.
|
||||
"""
|
||||
if use_contextual and self.bm25_contextual_index:
|
||||
index = self.bm25_contextual_index
|
||||
corpus = self.bm25_contextual_corpus
|
||||
chunk_store = self.contextual_chunk_store
|
||||
elif self.bm25_index:
|
||||
index = self.bm25_index
|
||||
corpus = self.bm25_corpus
|
||||
chunk_store = self.chunk_store
|
||||
else:
|
||||
logger.warning("BM25 index not available")
|
||||
return []
|
||||
|
||||
# Tokenize query (jieba 中文分词,兼容英文)
|
||||
query_tokens = _bm25_tokenize(query)
|
||||
|
||||
# Get BM25 scores
|
||||
scores = index.get_scores(query_tokens)
|
||||
|
||||
# Get top-k indices
|
||||
top_indices = np.argsort(scores)[-top_k:][::-1]
|
||||
|
||||
# Create results
|
||||
results = []
|
||||
chunk_list = list(chunk_store.values())
|
||||
|
||||
for idx in top_indices:
|
||||
if idx < len(chunk_list) and scores[idx] > 0:
|
||||
chunk = chunk_list[idx]
|
||||
result = ContextualSearchResult(
|
||||
doc_id=chunk.doc_id,
|
||||
chunk_id=chunk.chunk_id,
|
||||
text=chunk.text,
|
||||
score=float(scores[idx]),
|
||||
is_contextual=use_contextual,
|
||||
context_text=chunk.context if use_contextual else "",
|
||||
bm25_score=float(scores[idx]),
|
||||
retrieval_method="bm25",
|
||||
metadata={"method": "bm25", "contextual": use_contextual}
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
def _search_embeddings(self, query: str, use_contextual: bool, top_k: int) -> List[ContextualSearchResult]:
|
||||
"""
|
||||
Perform semantic search using embeddings.
|
||||
|
||||
Educational Note:
|
||||
Embedding search captures semantic meaning and relationships.
|
||||
Contextual chunks provide richer semantic information,
|
||||
helping find conceptually related content even without exact matches.
|
||||
"""
|
||||
try:
|
||||
# Use the retrieval pipeline for embedding search
|
||||
response = requests.post(
|
||||
f"{self.config.local_base_url}/search",
|
||||
json={
|
||||
"query": query,
|
||||
"mode": "embedding", # Use embedding mode
|
||||
"top_k": top_k,
|
||||
"filter": {"is_contextual": use_contextual} if self.enable_comparison else None
|
||||
}, timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
results = []
|
||||
data = response.json()
|
||||
|
||||
for item in data.get("results", []):
|
||||
# Map back to our chunk store
|
||||
chunk_id = item.get("metadata", {}).get("chunk_id", "")
|
||||
|
||||
if use_contextual and chunk_id in self.contextual_chunk_store:
|
||||
chunk = self.contextual_chunk_store[chunk_id]
|
||||
elif chunk_id in self.chunk_store:
|
||||
chunk = self.chunk_store[chunk_id]
|
||||
else:
|
||||
continue
|
||||
|
||||
result = ContextualSearchResult(
|
||||
doc_id=chunk.doc_id,
|
||||
chunk_id=chunk.chunk_id,
|
||||
text=chunk.text,
|
||||
score=item.get("score", 0.0),
|
||||
is_contextual=use_contextual,
|
||||
context_text=chunk.context if hasattr(chunk, 'context') else "",
|
||||
embedding_score=item.get("score", 0.0),
|
||||
retrieval_method="embedding",
|
||||
metadata={"method": "embedding", "contextual": use_contextual}
|
||||
)
|
||||
results.append(result)
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in embedding search: {e}")
|
||||
return []
|
||||
|
||||
def _rank_fusion(self,
|
||||
bm25_results: List[ContextualSearchResult],
|
||||
embedding_results: List[ContextualSearchResult],
|
||||
top_k: int) -> List[ContextualSearchResult]:
|
||||
"""
|
||||
Combine BM25 and embedding results using reciprocal rank fusion.
|
||||
|
||||
Educational Note:
|
||||
Rank fusion combines different retrieval signals:
|
||||
- BM25 provides strong exact matching
|
||||
- Embeddings provide semantic understanding
|
||||
- The combination often outperforms either method alone
|
||||
|
||||
We use Reciprocal Rank Fusion (RRF) which is simple but effective.
|
||||
"""
|
||||
fusion_scores = {}
|
||||
chunk_map = {}
|
||||
|
||||
# RRF constant (typically 60)
|
||||
k = 60
|
||||
|
||||
# Process BM25 results
|
||||
for rank, result in enumerate(bm25_results):
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
fusion_scores[result.chunk_id] = fusion_scores.get(result.chunk_id, 0) + rrf_score
|
||||
chunk_map[result.chunk_id] = result
|
||||
result.bm25_score = result.score
|
||||
|
||||
# Process embedding results
|
||||
for rank, result in enumerate(embedding_results):
|
||||
rrf_score = 1.0 / (k + rank + 1)
|
||||
|
||||
if result.chunk_id in fusion_scores:
|
||||
# Update existing result
|
||||
fusion_scores[result.chunk_id] += rrf_score
|
||||
chunk_map[result.chunk_id].embedding_score = result.score
|
||||
else:
|
||||
# New result from embeddings only
|
||||
fusion_scores[result.chunk_id] = rrf_score
|
||||
chunk_map[result.chunk_id] = result
|
||||
result.embedding_score = result.score
|
||||
|
||||
# Create final results sorted by fusion score
|
||||
final_results = []
|
||||
for chunk_id, fusion_score in sorted(fusion_scores.items(),
|
||||
key=lambda x: x[1],
|
||||
reverse=True)[:top_k]:
|
||||
result = chunk_map[chunk_id]
|
||||
result.hybrid_score = fusion_score
|
||||
result.score = fusion_score # Use fusion score as main score
|
||||
result.retrieval_method = "hybrid"
|
||||
final_results.append(result)
|
||||
|
||||
return final_results
|
||||
|
||||
def compare_retrieval_methods(self,
|
||||
query: str,
|
||||
top_k: int = 20) -> Dict[str, Any]:
|
||||
"""
|
||||
Compare contextual vs non-contextual retrieval.
|
||||
|
||||
Educational Note:
|
||||
This method demonstrates the improvement that contextual
|
||||
retrieval provides across different search methods.
|
||||
It's useful for evaluation and understanding when context helps most.
|
||||
"""
|
||||
logger.info(f"Comparing retrieval methods for: {query[:100]}...")
|
||||
|
||||
comparison_results = {
|
||||
"query": query,
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
"top_k": top_k,
|
||||
"methods": {}
|
||||
}
|
||||
|
||||
# Test each combination
|
||||
test_configs = [
|
||||
("contextual_hybrid", True, "hybrid"),
|
||||
("contextual_bm25", True, "bm25"),
|
||||
("contextual_embedding", True, "embedding"),
|
||||
("non_contextual_hybrid", False, "hybrid"),
|
||||
("non_contextual_bm25", False, "bm25"),
|
||||
("non_contextual_embedding", False, "embedding")
|
||||
]
|
||||
|
||||
for name, use_contextual, method in test_configs:
|
||||
# Temporarily set mode
|
||||
original_contextual = self.use_contextual
|
||||
self.use_contextual = use_contextual
|
||||
|
||||
# Perform search
|
||||
start_time = time.time()
|
||||
results = self.contextual_search(query, method, top_k)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Store results
|
||||
comparison_results["methods"][name] = {
|
||||
"results": [r.to_dict() for r in results[:5]], # Top 5 for readability
|
||||
"total_results": len(results),
|
||||
"retrieval_time": elapsed,
|
||||
"avg_score": np.mean([r.score for r in results]) if results else 0,
|
||||
"max_score": max([r.score for r in results]) if results else 0
|
||||
}
|
||||
|
||||
# Restore mode
|
||||
self.use_contextual = original_contextual
|
||||
|
||||
# Add analysis
|
||||
comparison_results["analysis"] = self._analyze_comparison(comparison_results)
|
||||
|
||||
# Update stats
|
||||
self.search_stats["comparison_searches"] += 1
|
||||
|
||||
return comparison_results
|
||||
|
||||
def _analyze_comparison(self, results: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Analyze comparison results to highlight improvements"""
|
||||
analysis = {
|
||||
"contextual_improvement": {},
|
||||
"method_comparison": {},
|
||||
"recommendations": []
|
||||
}
|
||||
|
||||
# Compare contextual vs non-contextual for each method
|
||||
for method in ["hybrid", "bm25", "embedding"]:
|
||||
contextual_key = f"contextual_{method}"
|
||||
non_contextual_key = f"non_contextual_{method}"
|
||||
|
||||
if contextual_key in results["methods"] and non_contextual_key in results["methods"]:
|
||||
contextual = results["methods"][contextual_key]
|
||||
non_contextual = results["methods"][non_contextual_key]
|
||||
|
||||
# Calculate improvement
|
||||
score_improvement = (
|
||||
(contextual["avg_score"] - non_contextual["avg_score"])
|
||||
/ non_contextual["avg_score"] * 100
|
||||
if non_contextual["avg_score"] > 0 else 0
|
||||
)
|
||||
|
||||
analysis["contextual_improvement"][method] = {
|
||||
"score_improvement_pct": round(score_improvement, 2),
|
||||
"contextual_avg_score": round(contextual["avg_score"], 4),
|
||||
"non_contextual_avg_score": round(non_contextual["avg_score"], 4)
|
||||
}
|
||||
|
||||
# Find best performing method
|
||||
best_method = max(
|
||||
results["methods"].items(),
|
||||
key=lambda x: x[1]["avg_score"]
|
||||
)
|
||||
analysis["best_method"] = best_method[0]
|
||||
|
||||
# Generate recommendations
|
||||
if "hybrid" in analysis["contextual_improvement"]:
|
||||
if analysis["contextual_improvement"]["hybrid"]["score_improvement_pct"] > 10:
|
||||
analysis["recommendations"].append(
|
||||
"Contextual retrieval shows significant improvement (>10%). "
|
||||
"Consider using it for production."
|
||||
)
|
||||
|
||||
if analysis["contextual_improvement"]["bm25"]["score_improvement_pct"] > \
|
||||
analysis["contextual_improvement"]["embedding"]["score_improvement_pct"]:
|
||||
analysis["recommendations"].append(
|
||||
"Contextual enhancement helps BM25 more than embeddings. "
|
||||
"The query might contain specific terms that benefit from context."
|
||||
)
|
||||
|
||||
return analysis
|
||||
|
||||
def _save_indexes(self):
|
||||
"""Save BM25 indexes to disk"""
|
||||
try:
|
||||
# Save BM25 indexes
|
||||
if self.bm25_contextual_index:
|
||||
with open(self.index_dir / "bm25_contextual.pkl", "wb") as f:
|
||||
pickle.dump({
|
||||
"index": self.bm25_contextual_index,
|
||||
"corpus": self.bm25_contextual_corpus
|
||||
}, f)
|
||||
|
||||
if self.bm25_index:
|
||||
with open(self.index_dir / "bm25_non_contextual.pkl", "wb") as f:
|
||||
pickle.dump({
|
||||
"index": self.bm25_index,
|
||||
"corpus": self.bm25_corpus
|
||||
}, f)
|
||||
|
||||
# Save chunk stores
|
||||
with open(self.index_dir / "chunk_stores.json", "w") as f:
|
||||
json.dump({
|
||||
"contextual": {k: v.to_dict() for k, v in self.contextual_chunk_store.items()},
|
||||
"non_contextual": {k: v.to_dict() for k, v in self.chunk_store.items()}
|
||||
}, f, indent=2)
|
||||
|
||||
logger.info("Indexes saved successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving indexes: {e}")
|
||||
|
||||
def _load_indexes(self):
|
||||
"""Load BM25 indexes from disk"""
|
||||
try:
|
||||
# Load BM25 indexes
|
||||
contextual_path = self.index_dir / "bm25_contextual.pkl"
|
||||
if contextual_path.exists():
|
||||
with open(contextual_path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
self.bm25_contextual_index = data["index"]
|
||||
self.bm25_contextual_corpus = data["corpus"]
|
||||
logger.info("Loaded contextual BM25 index")
|
||||
|
||||
non_contextual_path = self.index_dir / "bm25_non_contextual.pkl"
|
||||
if non_contextual_path.exists():
|
||||
with open(non_contextual_path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
self.bm25_index = data["index"]
|
||||
self.bm25_corpus = data["corpus"]
|
||||
logger.info("Loaded non-contextual BM25 index")
|
||||
|
||||
# Load chunk stores
|
||||
stores_path = self.index_dir / "chunk_stores.json"
|
||||
if stores_path.exists():
|
||||
with open(stores_path, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Reconstruct contextual chunks
|
||||
for chunk_id, chunk_dict in data.get("contextual", {}).items():
|
||||
self.contextual_chunk_store[chunk_id] = ContextualChunk(
|
||||
chunk_id=chunk_dict["chunk_id"],
|
||||
doc_id=chunk_dict["doc_id"],
|
||||
text=chunk_dict["text"],
|
||||
context=chunk_dict["context"],
|
||||
contextualized_text=chunk_dict["contextualized_text"],
|
||||
chunk_index=chunk_dict["chunk_index"],
|
||||
char_count=chunk_dict["char_count"],
|
||||
context_tokens=chunk_dict.get("context_tokens", 0),
|
||||
generation_time=chunk_dict.get("generation_time", 0),
|
||||
metadata=chunk_dict.get("metadata", {})
|
||||
)
|
||||
|
||||
# Reconstruct non-contextual chunks
|
||||
for chunk_id, chunk_dict in data.get("non_contextual", {}).items():
|
||||
self.chunk_store[chunk_id] = ContextualChunk(
|
||||
chunk_id=chunk_dict["chunk_id"],
|
||||
doc_id=chunk_dict["doc_id"],
|
||||
text=chunk_dict["text"],
|
||||
context="",
|
||||
contextualized_text=chunk_dict["text"],
|
||||
chunk_index=chunk_dict["chunk_index"],
|
||||
char_count=chunk_dict["char_count"],
|
||||
metadata=chunk_dict.get("metadata", {})
|
||||
)
|
||||
|
||||
logger.info(f"Loaded {len(self.contextual_chunk_store)} contextual chunks")
|
||||
logger.info(f"Loaded {len(self.chunk_store)} non-contextual chunks")
|
||||
|
||||
except Exception as e:
|
||||
logger.info(f"No existing indexes found or error loading: {e}")
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""Get comprehensive statistics"""
|
||||
stats = {
|
||||
"search_stats": self.search_stats,
|
||||
"index_stats": {
|
||||
"contextual_chunks": len(self.contextual_chunk_store),
|
||||
"non_contextual_chunks": len(self.chunk_store),
|
||||
"bm25_contextual_size": len(self.bm25_contextual_corpus) if self.bm25_contextual_corpus else 0,
|
||||
"bm25_non_contextual_size": len(self.bm25_corpus) if self.bm25_corpus else 0
|
||||
}
|
||||
}
|
||||
return stats
|
||||
@@ -0,0 +1,410 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Interactive demo of the Contextual Retrieval System
|
||||
|
||||
This script provides an interactive demonstration showing:
|
||||
1. How chunks lose context in traditional RAG
|
||||
2. How contextual retrieval solves this problem
|
||||
3. Side-by-side comparison of retrieval quality
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
from config import Config
|
||||
from contextual_chunking import ContextualChunker
|
||||
from contextual_tools import ContextualKnowledgeBaseTools
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(message)s' # Simple format for demo
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def print_header(title: str, char: str = "=", width: int = 80):
|
||||
"""Print a formatted header"""
|
||||
logger.info(f"\n{char * width}")
|
||||
logger.info(f"{title.center(width)}")
|
||||
logger.info(f"{char * width}\n")
|
||||
|
||||
|
||||
def print_section(title: str, char: str = "-", width: int = 60):
|
||||
"""Print a section header"""
|
||||
logger.info(f"\n{char * width}")
|
||||
logger.info(f"{title}")
|
||||
logger.info(f"{char * width}\n")
|
||||
|
||||
|
||||
class ContextualRetrievalDemo:
|
||||
"""Interactive demo class"""
|
||||
|
||||
def __init__(self):
|
||||
self.config = Config.from_env()
|
||||
self.documents = {}
|
||||
self.contextual_kb = None
|
||||
self.non_contextual_kb = None
|
||||
|
||||
def run(self):
|
||||
"""Run the interactive demo"""
|
||||
print_header("CONTEXTUAL RETRIEVAL SYSTEM - INTERACTIVE DEMO")
|
||||
|
||||
logger.info("Welcome! This demo will show you how contextual retrieval")
|
||||
logger.info("improves RAG systems by preserving context when chunking.\n")
|
||||
|
||||
while True:
|
||||
self.show_menu()
|
||||
choice = input("\nYour choice: ").strip()
|
||||
|
||||
if choice == "1":
|
||||
self.demo_problem()
|
||||
elif choice == "2":
|
||||
self.demo_solution()
|
||||
elif choice == "3":
|
||||
self.demo_comparison()
|
||||
elif choice == "4":
|
||||
self.demo_real_example()
|
||||
elif choice == "5":
|
||||
self.show_statistics()
|
||||
elif choice == "q":
|
||||
logger.info("\nThank you for using the Contextual Retrieval Demo!")
|
||||
break
|
||||
else:
|
||||
logger.info("Invalid choice. Please try again.")
|
||||
|
||||
def show_menu(self):
|
||||
"""Show main menu"""
|
||||
print_section("MAIN MENU")
|
||||
logger.info("1. Demonstrate the Context Loss Problem")
|
||||
logger.info("2. Show the Contextual Retrieval Solution")
|
||||
logger.info("3. Compare Search Results (Side-by-Side)")
|
||||
logger.info("4. Real-World Example (Financial Report)")
|
||||
logger.info("5. Show Performance Statistics")
|
||||
logger.info("q. Quit")
|
||||
|
||||
def demo_problem(self):
|
||||
"""Demonstrate the context loss problem"""
|
||||
print_header("THE CONTEXT LOSS PROBLEM", char="*")
|
||||
|
||||
# Example document
|
||||
document = """
|
||||
ACME Corporation Annual Report 2023
|
||||
|
||||
Financial Highlights:
|
||||
ACME Corporation achieved record performance in 2023. The company's revenue
|
||||
grew by 15% compared to the previous year, reaching $2.5 billion. This growth
|
||||
was driven by strong demand in the technology sector.
|
||||
|
||||
TechStart Inc. Performance:
|
||||
Meanwhile, TechStart Inc. faced challenges in 2023. The company's revenue
|
||||
declined by 8% due to supply chain disruptions. Management has implemented
|
||||
cost-cutting measures to improve profitability.
|
||||
|
||||
Global Industries Update:
|
||||
Global Industries maintained steady performance. The company's revenue
|
||||
remained flat at $1.8 billion, but profit margins improved by 2 percentage
|
||||
points through operational efficiency gains.
|
||||
""".strip()
|
||||
|
||||
logger.info("Consider this document:\n")
|
||||
logger.info("=" * 60)
|
||||
logger.info(document)
|
||||
logger.info("=" * 60)
|
||||
|
||||
logger.info("\nNow imagine we chunk this document and get:\n")
|
||||
|
||||
# Show problematic chunks
|
||||
chunks = [
|
||||
"The company's revenue grew by 15% compared to the previous year, reaching $2.5 billion.",
|
||||
"The company's revenue declined by 8% due to supply chain disruptions.",
|
||||
"The company's revenue remained flat at $1.8 billion, but profit margins improved."
|
||||
]
|
||||
|
||||
for i, chunk in enumerate(chunks, 1):
|
||||
logger.info(f"CHUNK {i}:")
|
||||
logger.info(f" '{chunk}'")
|
||||
logger.info("")
|
||||
|
||||
logger.info("❌ PROBLEM: All chunks say 'The company' but refer to different companies!")
|
||||
logger.info("❌ A search for 'company revenue growth' might return the wrong chunk!")
|
||||
logger.info("❌ Without context, we can't tell which company each chunk refers to!\n")
|
||||
|
||||
input("\nPress Enter to continue...")
|
||||
|
||||
def demo_solution(self):
|
||||
"""Demonstrate the contextual retrieval solution"""
|
||||
print_header("THE CONTEXTUAL RETRIEVAL SOLUTION", char="*")
|
||||
|
||||
logger.info("Contextual Retrieval solves this by adding context to each chunk:\n")
|
||||
|
||||
# Show contextualized chunks
|
||||
contextual_chunks = [
|
||||
{
|
||||
"context": "This chunk is from ACME Corporation's 2023 annual report financial highlights section.",
|
||||
"text": "The company's revenue grew by 15% compared to the previous year, reaching $2.5 billion."
|
||||
},
|
||||
{
|
||||
"context": "This chunk discusses TechStart Inc.'s 2023 performance challenges.",
|
||||
"text": "The company's revenue declined by 8% due to supply chain disruptions."
|
||||
},
|
||||
{
|
||||
"context": "This chunk covers Global Industries' steady 2023 performance.",
|
||||
"text": "The company's revenue remained flat at $1.8 billion, but profit margins improved."
|
||||
}
|
||||
]
|
||||
|
||||
for i, chunk in enumerate(contextual_chunks, 1):
|
||||
logger.info(f"CONTEXTUAL CHUNK {i}:")
|
||||
logger.info(f" Context: {chunk['context']}")
|
||||
logger.info(f" Text: {chunk['text']}")
|
||||
logger.info(f" Combined: {chunk['context']} {chunk['text']}\n")
|
||||
|
||||
logger.info("✅ SOLUTION: Each chunk now has context!")
|
||||
logger.info("✅ Searching for 'ACME revenue growth' will correctly find chunk 1!")
|
||||
logger.info("✅ The context preserves crucial information lost in traditional chunking!\n")
|
||||
|
||||
input("\nPress Enter to continue...")
|
||||
|
||||
def demo_comparison(self):
|
||||
"""Run a side-by-side comparison"""
|
||||
print_header("SIDE-BY-SIDE COMPARISON", char="*")
|
||||
|
||||
# Create test document
|
||||
test_doc = """
|
||||
Artificial Intelligence in Healthcare
|
||||
|
||||
Introduction:
|
||||
Artificial intelligence is transforming healthcare delivery. Machine learning
|
||||
models are being used for disease diagnosis, drug discovery, and patient care
|
||||
optimization. The technology has shown remarkable results in early detection
|
||||
of diseases.
|
||||
|
||||
Diagnostic Applications:
|
||||
In radiology, AI systems can detect cancer with 95% accuracy. The systems
|
||||
analyze medical images faster than human radiologists. This reduces diagnosis
|
||||
time from hours to minutes.
|
||||
|
||||
Drug Discovery:
|
||||
Pharmaceutical companies use AI to identify potential drug compounds. The
|
||||
technology can predict drug interactions and side effects. This accelerates
|
||||
the drug development process by years.
|
||||
""".strip()
|
||||
|
||||
logger.info("Test Document:")
|
||||
logger.info("=" * 60)
|
||||
logger.info(test_doc[:300] + "..." if len(test_doc) > 300 else test_doc)
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Initialize systems
|
||||
logger.info("\nInitializing systems...")
|
||||
|
||||
# Create contextual system
|
||||
contextual_chunker = ContextualChunker(
|
||||
chunking_config=self.config.chunking,
|
||||
llm_config=self.config.llm,
|
||||
use_contextual=True
|
||||
)
|
||||
self.contextual_kb = ContextualKnowledgeBaseTools(
|
||||
config=self.config.knowledge_base,
|
||||
use_contextual=True
|
||||
)
|
||||
|
||||
# Create non-contextual system
|
||||
non_contextual_chunker = ContextualChunker(
|
||||
chunking_config=self.config.chunking,
|
||||
llm_config=self.config.llm,
|
||||
use_contextual=False
|
||||
)
|
||||
self.non_contextual_kb = ContextualKnowledgeBaseTools(
|
||||
config=self.config.knowledge_base,
|
||||
use_contextual=False
|
||||
)
|
||||
|
||||
# Process document
|
||||
logger.info("\nProcessing document...")
|
||||
|
||||
# Contextual chunks
|
||||
contextual_chunks = contextual_chunker.chunk_document(
|
||||
text=test_doc,
|
||||
doc_id="healthcare_ai"
|
||||
)
|
||||
self.contextual_kb.index_contextual_chunks(contextual_chunks)
|
||||
|
||||
# Non-contextual chunks
|
||||
non_contextual_chunks = non_contextual_chunker.chunk_document(
|
||||
text=test_doc,
|
||||
doc_id="healthcare_ai"
|
||||
)
|
||||
self.non_contextual_kb.index_contextual_chunks(non_contextual_chunks)
|
||||
|
||||
# Test queries
|
||||
test_queries = [
|
||||
"How accurate is the AI system?",
|
||||
"What technology reduces diagnosis time?",
|
||||
"What can the technology predict?"
|
||||
]
|
||||
|
||||
logger.info("\nRunning comparison...")
|
||||
|
||||
for query in test_queries:
|
||||
print_section(f"Query: {query}")
|
||||
|
||||
# Contextual search
|
||||
contextual_results = self.contextual_kb.contextual_search(
|
||||
query=query,
|
||||
method="hybrid",
|
||||
top_k=3
|
||||
)
|
||||
|
||||
# Non-contextual search
|
||||
non_contextual_results = self.non_contextual_kb.contextual_search(
|
||||
query=query,
|
||||
method="hybrid",
|
||||
top_k=3
|
||||
)
|
||||
|
||||
logger.info("CONTEXTUAL RESULTS:")
|
||||
if contextual_results:
|
||||
result = contextual_results[0]
|
||||
logger.info(f" Score: {result.score:.4f}")
|
||||
logger.info(f" Context: {result.context_text[:100]}..." if result.context_text else " Context: None")
|
||||
logger.info(f" Match: {result.text[:100]}...\n")
|
||||
else:
|
||||
logger.info(" No results found\n")
|
||||
|
||||
logger.info("NON-CONTEXTUAL RESULTS:")
|
||||
if non_contextual_results:
|
||||
result = non_contextual_results[0]
|
||||
logger.info(f" Score: {result.score:.4f}")
|
||||
logger.info(f" Match: {result.text[:100]}...\n")
|
||||
else:
|
||||
logger.info(" No results found\n")
|
||||
|
||||
# Score comparison
|
||||
if contextual_results and non_contextual_results:
|
||||
improvement = ((contextual_results[0].score - non_contextual_results[0].score)
|
||||
/ non_contextual_results[0].score * 100)
|
||||
logger.info(f"📊 Score Improvement: {improvement:.1f}%\n")
|
||||
|
||||
input("\nPress Enter to continue...")
|
||||
|
||||
def demo_real_example(self):
|
||||
"""Demonstrate with a real-world example"""
|
||||
print_header("REAL-WORLD EXAMPLE: FINANCIAL REPORT", char="*")
|
||||
|
||||
# Create a realistic financial report
|
||||
report = """
|
||||
Q2 2023 Earnings Report - TechCorp International
|
||||
|
||||
Executive Summary:
|
||||
TechCorp International reported strong second quarter results for 2023,
|
||||
with revenue of $850 million, representing a 12% year-over-year growth.
|
||||
The company's cloud services division was the primary growth driver.
|
||||
|
||||
Revenue Breakdown:
|
||||
Cloud Services: Revenue increased by 25% to $400 million, driven by
|
||||
enterprise adoption of our AI-powered analytics platform. Operating margin
|
||||
improved to 35% from 30% in the prior year.
|
||||
|
||||
Hardware Division: Revenue declined by 5% to $300 million due to supply
|
||||
chain constraints. However, the new product pipeline remains strong with
|
||||
three launches planned for Q3.
|
||||
|
||||
Software Licensing: Revenue grew by 8% to $150 million. The company added
|
||||
200 new enterprise customers during the quarter, bringing the total to
|
||||
5,000 active licenses.
|
||||
|
||||
Competitive Analysis:
|
||||
Compared to DataSoft Corp, our main competitor, we maintained market share
|
||||
leadership. DataSoft reported 8% revenue growth in their latest quarter,
|
||||
while our 12% growth demonstrates strong execution. Their cloud division
|
||||
grew 15% compared to our 25% growth.
|
||||
|
||||
Future Outlook:
|
||||
Management expects continued momentum in Q3 2023. The company raised
|
||||
full-year guidance to $3.5 billion in revenue, representing 15% annual
|
||||
growth. Investment in R&D will increase by 20% to accelerate AI
|
||||
product development.
|
||||
""".strip()
|
||||
|
||||
logger.info("Processing a realistic financial report...")
|
||||
logger.info("=" * 60)
|
||||
logger.info(report[:400] + "..." if len(report) > 400 else report)
|
||||
logger.info("=" * 60)
|
||||
|
||||
# Process with contextual system
|
||||
logger.info("\nGenerating contextual chunks (this may take a moment)...\n")
|
||||
|
||||
chunker = ContextualChunker(
|
||||
chunking_config=self.config.chunking,
|
||||
llm_config=self.config.llm,
|
||||
use_contextual=True
|
||||
)
|
||||
|
||||
chunks = chunker.chunk_document(
|
||||
text=report,
|
||||
doc_id="techcorp_q2_2023"
|
||||
)
|
||||
|
||||
# Show examples of contextualized chunks
|
||||
logger.info("Example Contextual Chunks:\n")
|
||||
|
||||
for chunk in chunks[:3]:
|
||||
logger.info(f"Original: {chunk.text[:100]}...")
|
||||
logger.info(f"Context: {chunk.context}")
|
||||
logger.info("")
|
||||
|
||||
# Show how this helps with ambiguous queries
|
||||
logger.info("Why this matters:\n")
|
||||
logger.info("Query: 'What was the revenue growth?'")
|
||||
logger.info(" - Without context: Could match TechCorp's 12%, Cloud's 25%, or Software's 8%")
|
||||
logger.info(" - With context: Correctly identifies which growth figure you want\n")
|
||||
|
||||
logger.info("Query: 'How did the company perform vs competition?'")
|
||||
logger.info(" - Without context: Might return DataSoft's results")
|
||||
logger.info(" - With context: Returns TechCorp's performance comparison\n")
|
||||
|
||||
input("\nPress Enter to continue...")
|
||||
|
||||
def show_statistics(self):
|
||||
"""Show performance statistics"""
|
||||
print_header("PERFORMANCE STATISTICS", char="*")
|
||||
|
||||
if not self.contextual_kb:
|
||||
logger.info("No searches performed yet. Run a comparison first!\n")
|
||||
input("\nPress Enter to continue...")
|
||||
return
|
||||
|
||||
# Get statistics
|
||||
stats = self.contextual_kb.get_statistics()
|
||||
|
||||
logger.info("Search Statistics:")
|
||||
logger.info(f" Total searches: {stats['search_stats']['total_searches']}")
|
||||
logger.info(f" Average retrieval time: {stats['search_stats']['avg_retrieval_time']:.3f}s")
|
||||
logger.info(f" Contextual searches: {stats['search_stats']['contextual_searches']}")
|
||||
logger.info(f" Non-contextual searches: {stats['search_stats']['non_contextual_searches']}")
|
||||
|
||||
logger.info("\nIndex Statistics:")
|
||||
logger.info(f" Contextual chunks indexed: {stats['index_stats']['contextual_chunks']}")
|
||||
logger.info(f" Non-contextual chunks indexed: {stats['index_stats']['non_contextual_chunks']}")
|
||||
|
||||
logger.info("\nBased on Anthropic's Research:")
|
||||
logger.info(" Standard RAG: 5.7% retrieval failure rate")
|
||||
logger.info(" Contextual RAG: 2.9% failure rate (49% improvement)")
|
||||
logger.info(" + Reranking: 1.9% failure rate (67% improvement)")
|
||||
|
||||
input("\nPress Enter to continue...")
|
||||
|
||||
|
||||
def main():
|
||||
"""Run the interactive demo"""
|
||||
demo = ContextualRetrievalDemo()
|
||||
demo.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,71 @@
|
||||
# Contextual Retrieval System Configuration
|
||||
# Copy this file to .env and fill in your API keys
|
||||
|
||||
# LLM Provider Configuration
|
||||
# Choose your preferred provider for context generation
|
||||
|
||||
# Kimi/Moonshot (Recommended for Chinese content)
|
||||
MOONSHOT_API_KEY=your_kimi_api_key_here
|
||||
|
||||
# Doubao
|
||||
ARK_API_KEY=your_doubao_api_key_here
|
||||
|
||||
# Alibaba Cloud Model Studio / Bailian (Qwen)
|
||||
DASHSCOPE_API_KEY=your_dashscope_api_key_here
|
||||
# DASHSCOPE_BASE_URL=https://dashscope-intl.aliyuncs.com/compatible-mode/v1
|
||||
|
||||
# OpenAI (Good for general use)
|
||||
OPENAI_API_KEY=your_openai_api_key_here
|
||||
|
||||
# SiliconFlow
|
||||
SILICONFLOW_API_KEY=your_siliconflow_api_key_here
|
||||
|
||||
# OpenRouter: usable as an explicit LLM_PROVIDER, and also a universal fallback
|
||||
# — if the configured provider's key is missing but OPENROUTER_API_KEY is set,
|
||||
# the agent auto-routes through OpenRouter (model names mapped automatically;
|
||||
# set OPENROUTER_MODEL to override).
|
||||
OPENROUTER_API_KEY=your_openrouter_api_key_here
|
||||
|
||||
# Groq (Fast inference)
|
||||
GROQ_API_KEY=your_groq_api_key_here
|
||||
|
||||
# Together AI
|
||||
TOGETHER_API_KEY=your_together_api_key_here
|
||||
|
||||
# DeepSeek
|
||||
DEEPSEEK_API_KEY=your_deepseek_api_key_here
|
||||
|
||||
# Default LLM Configuration
|
||||
LLM_PROVIDER=openai # Options include dashscope/qwen/bailian, kimi, doubao, openai, siliconflow, etc.
|
||||
LLM_MODEL=gpt-5.6-luna # Model for context generation (use cheaper models to save cost)
|
||||
LLM_TEMPERATURE=0.3 # Lower temperature for consistent context generation
|
||||
LLM_MAX_TOKENS=150 # Max tokens for context generation (2-3 sentences)
|
||||
|
||||
# Knowledge Base Configuration
|
||||
KB_TYPE=local # Options: local, dify, raptor, graphrag
|
||||
KB_LOCAL_BASE_URL=http://localhost:4242 # Local retrieval pipeline URL
|
||||
|
||||
# Chunking Configuration
|
||||
CHUNK_SIZE=2048 # Characters per chunk
|
||||
MAX_CHUNK_SIZE=1024 # Maximum chunk size
|
||||
CHUNK_OVERLAP=200 # Overlap between chunks
|
||||
RESPECT_PARAGRAPH_BOUNDARY=true # Preserve paragraph structure
|
||||
|
||||
# Contextual Retrieval Settings
|
||||
USE_CONTEXTUAL=true # Enable contextual retrieval by default
|
||||
ENABLE_COMPARISON=false # Enable comparison mode for evaluation
|
||||
CACHE_CONTEXTS=true # Cache generated contexts to reduce API calls
|
||||
|
||||
# Agent Configuration
|
||||
AGENT_MAX_ITERATIONS=10 # Max iterations for ReAct loop
|
||||
AGENT_VERBOSE=true # Enable detailed logging
|
||||
CONVERSATION_HISTORY_LIMIT=20 # Number of messages to keep in history
|
||||
|
||||
# Performance Settings
|
||||
BATCH_SIZE=10 # Number of chunks to process in parallel
|
||||
MAX_WORKERS=4 # Maximum parallel workers for processing
|
||||
CACHE_SIZE=1000 # Maximum number of contexts to cache
|
||||
|
||||
# Cost Control
|
||||
MAX_CONTEXT_GENERATION_COST=10.0 # Maximum cost in USD for context generation
|
||||
WARN_AT_COST=5.0 # Warn when cost exceeds this amount
|
||||
@@ -0,0 +1,317 @@
|
||||
"""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")
|
||||
@@ -0,0 +1,396 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,96 @@
|
||||
{
|
||||
"description": "上下文感知检索 vs. 传统分块的检索召回评测集(对应实验 3-10)。gold_chunk_id 为人工标注的相关文本块,标注依据见 note 字段。语料为 document_store.json 中已建好索引的《宪法》与《检察官法》分块。",
|
||||
"corpus": "document_store.json",
|
||||
"queries": [
|
||||
{
|
||||
"id": "q01",
|
||||
"query": "宪法是哪一年通过的?后来经过了几次修正?",
|
||||
"gold_chunk_id": "宪法_chunk_0",
|
||||
"note": "chunk_0 为宪法开篇,列举 1982 年通过日期及历次修正案日期。"
|
||||
},
|
||||
{
|
||||
"id": "q02",
|
||||
"query": "国家的根本任务是什么?",
|
||||
"gold_chunk_id": "宪法_chunk_1",
|
||||
"note": "chunk_1 序言明确写明“国家的根本任务是……”。"
|
||||
},
|
||||
{
|
||||
"id": "q03",
|
||||
"query": "我国处理民族关系的基本原则和对外政策中的五项原则是什么?",
|
||||
"gold_chunk_id": "宪法_chunk_2",
|
||||
"note": "chunk_2 阐述社会主义民族关系与和平共处五项原则。"
|
||||
},
|
||||
{
|
||||
"id": "q04",
|
||||
"query": "少数民族聚居的地方如何实行区域自治?",
|
||||
"gold_chunk_id": "宪法_chunk_3",
|
||||
"note": "chunk_3 第四条前后规定民族区域自治与依法治国。"
|
||||
},
|
||||
{
|
||||
"id": "q05",
|
||||
"query": "宪法如何保护个体经济、私营经济等非公有制经济?",
|
||||
"gold_chunk_id": "宪法_chunk_4",
|
||||
"note": "chunk_4 含第十一至十三条,规定非公有制经济与私有财产保护。"
|
||||
},
|
||||
{
|
||||
"id": "q06",
|
||||
"query": "宪法关于设立特别行政区以及在华外国人权利的规定是什么?",
|
||||
"gold_chunk_id": "宪法_chunk_6",
|
||||
"note": "chunk_6 含第三十一、三十二条,特别行政区与外国人条款。"
|
||||
},
|
||||
{
|
||||
"id": "q07",
|
||||
"query": "公民对国家机关工作人员的违法失职行为享有哪些监督权利?",
|
||||
"gold_chunk_id": "宪法_chunk_7",
|
||||
"note": "chunk_7 第四十一条规定批评、建议、申诉、控告、检举权。"
|
||||
},
|
||||
{
|
||||
"id": "q08",
|
||||
"query": "全国人民代表大会常务委员会有哪些职权?",
|
||||
"gold_chunk_id": "宪法_chunk_10",
|
||||
"note": "旗舰对照案例:chunk_10 原文以“(四)解释法律……”开头,正文既无“常务委员会”也无“职权”字样,纯靠上下文前缀(“第六十七条 全国人民代表大会常务委员会行使下列职权”)才能被锚定;无上下文时极易被 chunk_9/chunk_11 抢占。"
|
||||
},
|
||||
{
|
||||
"id": "q09",
|
||||
"query": "国家主席有哪些职权?",
|
||||
"gold_chunk_id": "宪法_chunk_12",
|
||||
"note": "chunk_12 含第八十至八十二条,国家主席公布法律、任免国务院总理等职权。"
|
||||
},
|
||||
{
|
||||
"id": "q10",
|
||||
"query": "地方各级人民代表大会的任期和代表选举方式是怎样规定的?",
|
||||
"gold_chunk_id": "宪法_chunk_14",
|
||||
"note": "chunk_14 第九十七至九十九条规定地方人大代表选举与五年任期。"
|
||||
},
|
||||
{
|
||||
"id": "q11",
|
||||
"query": "县级以上地方各级人民代表大会常务委员会组成人员有什么兼职限制?",
|
||||
"gold_chunk_id": "宪法_chunk_15",
|
||||
"note": "chunk_15 规定常委会组成人员不得担任行政、监察、审判、检察机关职务。"
|
||||
},
|
||||
{
|
||||
"id": "q12",
|
||||
"query": "监察委员会的性质、组成和任期是怎样规定的?",
|
||||
"gold_chunk_id": "宪法_chunk_17",
|
||||
"note": "chunk_17 第七节监察委员会,规定其性质、组成与任期。"
|
||||
},
|
||||
{
|
||||
"id": "q13",
|
||||
"query": "检察官法是什么时候修订的?",
|
||||
"gold_chunk_id": "检察官法_2019_04_23_chunk_0",
|
||||
"note": "chunk_0 列举检察官法 1995 年通过及历次修正、2019 年修订日期。"
|
||||
},
|
||||
{
|
||||
"id": "q14",
|
||||
"query": "检察官遴选委员会由哪些人员组成?律师参加公开选拔需要什么条件?",
|
||||
"gold_chunk_id": "检察官法_2019_04_23_chunk_2",
|
||||
"note": "chunk_2 第十六条规定省级检察官遴选委员会组成及律师公开选拔条件。"
|
||||
},
|
||||
{
|
||||
"id": "q15",
|
||||
"query": "检察官在哪些情形下应当予以免职?",
|
||||
"gold_chunk_id": "检察官法_2019_04_23_chunk_3",
|
||||
"note": "chunk_3 列举检察官免职情形及违反条件任命的撤销程序。"
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,588 @@
|
||||
"""Script to chunk and index local legal documents using Contextual Retrieval
|
||||
|
||||
This script:
|
||||
1. Cleans up existing indexes
|
||||
2. Reads legal documents from local laws directory
|
||||
3. Chunks them with paragraph-aware boundaries (soft limit 1024, hard limit 2048)
|
||||
4. Generates contextual descriptions for each chunk
|
||||
5. Indexes them with contextual enhancement
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Dict, Any, Optional
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass
|
||||
import time
|
||||
import requests
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.append(str(Path(__file__).parent))
|
||||
|
||||
from config import Config, ChunkingConfig, LLMConfig
|
||||
from contextual_chunking import ContextualChunker, ContextualChunk
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Configuration
|
||||
RETRIEVAL_PIPELINE_URL = "http://localhost:4242" # Default retrieval pipeline URL
|
||||
LAWS_DIR = Path("laws") # Local laws directory
|
||||
|
||||
# Custom chunking configuration for legal documents
|
||||
LEGAL_CHUNKING_CONFIG = ChunkingConfig(
|
||||
chunk_size=1024, # Soft limit
|
||||
max_chunk_size=2048, # Hard limit
|
||||
chunk_overlap=200,
|
||||
respect_paragraph_boundary=True,
|
||||
min_chunk_size=500
|
||||
)
|
||||
|
||||
|
||||
class ContextualLegalIndexer:
|
||||
"""Handles contextual chunking and indexing of local legal documents"""
|
||||
|
||||
def __init__(self,
|
||||
laws_dir: Path = LAWS_DIR,
|
||||
pipeline_url: str = RETRIEVAL_PIPELINE_URL,
|
||||
use_contextual: bool = True,
|
||||
llm_config: Optional[LLMConfig] = None):
|
||||
self.laws_dir = laws_dir
|
||||
self.pipeline_url = pipeline_url
|
||||
self.use_contextual = use_contextual
|
||||
|
||||
# Initialize contextual chunker
|
||||
self.chunker = ContextualChunker(
|
||||
chunking_config=LEGAL_CHUNKING_CONFIG,
|
||||
llm_config=llm_config or LLMConfig(),
|
||||
use_contextual=use_contextual
|
||||
)
|
||||
|
||||
self.stats = {
|
||||
"documents_processed": 0,
|
||||
"chunks_created": 0,
|
||||
"contextual_chunks": 0,
|
||||
"chunks_indexed": 0,
|
||||
"total_context_tokens": 0,
|
||||
"errors": 0,
|
||||
"categories_processed": set()
|
||||
}
|
||||
|
||||
# Document store for tracking (must match tools.py expectation)
|
||||
self.doc_store_path = Path("document_store.json")
|
||||
|
||||
logger.info(f"Initialized contextual indexer for local laws in {laws_dir}")
|
||||
logger.info(f"Pipeline URL: {pipeline_url}")
|
||||
logger.info(f"Contextual mode: {use_contextual}")
|
||||
|
||||
def cleanup_existing_index(self):
|
||||
"""Clean up existing indexes and document store"""
|
||||
logger.info("Cleaning up existing indexes...")
|
||||
|
||||
# Clean local document store
|
||||
if self.doc_store_path.exists():
|
||||
try:
|
||||
# Load existing store to get document IDs
|
||||
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
|
||||
existing_docs = json.load(f)
|
||||
|
||||
logger.info(f"Found {len(existing_docs)} existing documents in store")
|
||||
|
||||
# Clear the store
|
||||
self.doc_store_path.unlink()
|
||||
logger.info("Cleared local document store")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning document store: {e}")
|
||||
|
||||
# Try to clear the retrieval pipeline
|
||||
try:
|
||||
response = requests.delete(f"{self.pipeline_url}/clear", timeout=30)
|
||||
if response.status_code == 200:
|
||||
logger.info("Cleared retrieval pipeline index")
|
||||
else:
|
||||
logger.warning(f"Failed to clear pipeline: {response.status_code}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not clear retrieval pipeline: {e}")
|
||||
|
||||
logger.info("Cleanup complete")
|
||||
|
||||
def get_all_legal_documents(self) -> List[Dict[str, Any]]:
|
||||
"""Get all legal documents from local laws directory"""
|
||||
documents = []
|
||||
|
||||
if not self.laws_dir.exists():
|
||||
logger.error(f"Laws directory not found: {self.laws_dir}")
|
||||
return documents
|
||||
|
||||
# Iterate through category directories
|
||||
for category_dir in sorted(self.laws_dir.iterdir()):
|
||||
if not category_dir.is_dir():
|
||||
continue
|
||||
|
||||
category_name = category_dir.name
|
||||
logger.info(f"Found category: {category_name}")
|
||||
|
||||
# Find all .md files in this category
|
||||
for md_file in category_dir.glob("*.md"):
|
||||
doc_info = {
|
||||
"path": md_file,
|
||||
"name": md_file.stem, # filename without extension
|
||||
"category": category_name,
|
||||
"full_name": md_file.name
|
||||
}
|
||||
documents.append(doc_info)
|
||||
|
||||
logger.info(f"Found {len(documents)} legal documents")
|
||||
return documents
|
||||
|
||||
def read_document(self, doc_info: Dict[str, Any]) -> Optional[str]:
|
||||
"""Read a legal document from disk"""
|
||||
try:
|
||||
doc_path = doc_info["path"]
|
||||
content = doc_path.read_text(encoding='utf-8')
|
||||
logger.debug(f"Read {doc_info['name']} ({len(content)} chars)")
|
||||
return content
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading {doc_info['name']}: {e}")
|
||||
self.stats["errors"] += 1
|
||||
return None
|
||||
|
||||
def generate_document_id(self, doc_info: Dict[str, Any]) -> str:
|
||||
"""
|
||||
Generate a semantically meaningful document ID from file name.
|
||||
|
||||
Examples:
|
||||
- "宪法.md" -> "宪法"
|
||||
- "劳动法(2018-12-29).md" -> "劳动法_2018-12-29"
|
||||
- "民法典/第一编_总则.md" -> "民法典_第一编_总则"
|
||||
"""
|
||||
# Use the file stem (name without extension)
|
||||
base_name = doc_info["name"]
|
||||
|
||||
# Clean up the name: remove problematic characters but keep Chinese and alphanumeric
|
||||
import re
|
||||
# Replace parentheses with underscore
|
||||
clean_name = base_name.replace('(', '_').replace(')', '')
|
||||
clean_name = clean_name.replace('(', '_').replace(')', '')
|
||||
# Replace spaces and other separators with underscore
|
||||
clean_name = re.sub(r'[\s\-]+', '_', clean_name)
|
||||
# Remove trailing underscores
|
||||
clean_name = clean_name.strip('_')
|
||||
|
||||
# Don't add category prefix - just use the clean file name
|
||||
# This makes the IDs cleaner and more readable
|
||||
doc_id = clean_name
|
||||
|
||||
# Ensure the ID is not too long
|
||||
if len(doc_id) > 100:
|
||||
# Truncate but keep it meaningful
|
||||
doc_id = doc_id[:97] + "..."
|
||||
|
||||
logger.debug(f"Generated document ID: {doc_info['name']} -> {doc_id}")
|
||||
return doc_id
|
||||
|
||||
def process_document(self, doc_info: Dict[str, Any], content: str,
|
||||
index_immediately: bool = True) -> List[ContextualChunk]:
|
||||
"""Process a document using contextual chunking with optional immediate indexing"""
|
||||
# Generate semantically meaningful document ID
|
||||
doc_id = self.generate_document_id(doc_info)
|
||||
|
||||
# Prepare document metadata
|
||||
doc_metadata = {
|
||||
"category": doc_info["category"],
|
||||
"file_name": doc_info["full_name"],
|
||||
"document_type": "legal",
|
||||
"language": "zh-CN"
|
||||
}
|
||||
|
||||
# Create callback for immediate indexing if enabled
|
||||
indexed_chunks = []
|
||||
|
||||
def on_chunk_ready(chunk: ContextualChunk):
|
||||
"""Callback to index chunk immediately when ready"""
|
||||
# Update metadata
|
||||
chunk.metadata["category"] = doc_info["category"]
|
||||
chunk.metadata["doc_title"] = doc_info["name"]
|
||||
|
||||
if index_immediately:
|
||||
# Index the chunk immediately
|
||||
success = self.index_chunk(chunk)
|
||||
if success:
|
||||
logger.info(f" → Indexed chunk {chunk.chunk_index + 1} immediately")
|
||||
indexed_chunks.append(chunk)
|
||||
else:
|
||||
logger.warning(f" → Failed to index chunk {chunk.chunk_index + 1}")
|
||||
|
||||
# Use the contextual chunker with callback
|
||||
logger.info(f"Chunking {doc_info['name']} with contextual enhancement...")
|
||||
chunks = self.chunker.chunk_document(
|
||||
text=content,
|
||||
doc_id=doc_id,
|
||||
doc_metadata=doc_metadata,
|
||||
on_chunk_ready=on_chunk_ready if index_immediately else None
|
||||
)
|
||||
|
||||
# If not indexing immediately, update metadata for all chunks
|
||||
if not index_immediately:
|
||||
for chunk in chunks:
|
||||
chunk.metadata["category"] = doc_info["category"]
|
||||
chunk.metadata["doc_title"] = doc_info["name"]
|
||||
|
||||
return chunks
|
||||
|
||||
def index_chunk(self, chunk: ContextualChunk) -> bool:
|
||||
"""Index a contextual chunk in the retrieval pipeline"""
|
||||
try:
|
||||
# Prepare the indexing request
|
||||
# Use contextualized text for better retrieval
|
||||
index_data = {
|
||||
"text": chunk.contextualized_text if self.use_contextual else chunk.text,
|
||||
"doc_id": chunk.chunk_id,
|
||||
"metadata": {
|
||||
**chunk.metadata,
|
||||
"original_text": chunk.text,
|
||||
"context": chunk.context,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"char_count": chunk.char_count,
|
||||
"contextual": self.use_contextual
|
||||
}
|
||||
}
|
||||
|
||||
# Send to retrieval pipeline
|
||||
response = requests.post(
|
||||
f"{self.pipeline_url}/index",
|
||||
json=index_data,
|
||||
headers={"Content-Type": "application/json"}, timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
self.stats["chunks_indexed"] += 1
|
||||
if chunk.context:
|
||||
self.stats["contextual_chunks"] += 1
|
||||
|
||||
# IMPORTANT: Save to document store immediately after successful indexing
|
||||
self.save_chunk_to_document_store(chunk)
|
||||
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"Failed to index chunk {chunk.chunk_id}: {response.status_code}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error indexing chunk {chunk.chunk_id}: {e}")
|
||||
self.stats["errors"] += 1
|
||||
return False
|
||||
|
||||
def save_chunk_to_document_store(self, chunk: ContextualChunk):
|
||||
"""Save individual chunk to document store immediately (compatible with tools.py)"""
|
||||
try:
|
||||
# Load existing store or create new
|
||||
if self.doc_store_path.exists():
|
||||
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
|
||||
doc_store = json.load(f)
|
||||
else:
|
||||
doc_store = {}
|
||||
|
||||
# Save chunk in format expected by tools.py get_document method
|
||||
doc_store[chunk.chunk_id] = {
|
||||
"doc_id": chunk.chunk_id,
|
||||
"content": chunk.contextualized_text if self.use_contextual else chunk.text,
|
||||
"metadata": {
|
||||
**chunk.metadata,
|
||||
"original_text": chunk.text,
|
||||
"context": chunk.context,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"char_count": chunk.char_count,
|
||||
"contextual": self.use_contextual
|
||||
}
|
||||
}
|
||||
|
||||
# Write immediately - this is critical per user requirement
|
||||
with open(self.doc_store_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(doc_store, f, ensure_ascii=False, indent=2)
|
||||
|
||||
logger.debug(f"Saved chunk {chunk.chunk_id} to document store")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving chunk {chunk.chunk_id} to document store: {e}")
|
||||
|
||||
def save_document_info(self, doc_info: Dict[str, Any], chunks: List[ContextualChunk]):
|
||||
"""Save document summary information to local store"""
|
||||
# Load existing store or create new
|
||||
if self.doc_store_path.exists():
|
||||
with open(self.doc_store_path, 'r', encoding='utf-8') as f:
|
||||
doc_store = json.load(f)
|
||||
else:
|
||||
doc_store = {}
|
||||
|
||||
# Add document info with meaningful ID
|
||||
doc_id = self.generate_document_id(doc_info)
|
||||
|
||||
# Calculate statistics
|
||||
total_context_tokens = sum(c.context_tokens for c in chunks)
|
||||
avg_context_tokens = total_context_tokens / len(chunks) if chunks else 0
|
||||
|
||||
# Store document-level summary (compatible with tools.py format)
|
||||
doc_store[doc_id] = {
|
||||
"doc_id": doc_id,
|
||||
"content": f"Document: {doc_info['name']}\nCategory: {doc_info['category']}\n\nThis is a summary entry for the complete document. Individual chunks are stored separately.",
|
||||
"metadata": {
|
||||
"title": doc_info["name"],
|
||||
"category": doc_info["category"],
|
||||
"file": str(doc_info["path"]),
|
||||
"chunks": len(chunks),
|
||||
"contextual_chunks": sum(1 for c in chunks if c.context),
|
||||
"total_chars": sum(c.char_count for c in chunks),
|
||||
"total_context_tokens": total_context_tokens,
|
||||
"avg_context_tokens": avg_context_tokens,
|
||||
"indexed_at": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"chunk_ids": [c.chunk_id for c in chunks],
|
||||
"is_summary": True
|
||||
}
|
||||
}
|
||||
|
||||
# Save store
|
||||
with open(self.doc_store_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(doc_store, f, ensure_ascii=False, indent=2)
|
||||
|
||||
def process_all_documents(self,
|
||||
max_docs: Optional[int] = None,
|
||||
categories: Optional[List[str]] = None,
|
||||
batch_size: int = 10):
|
||||
"""Process all legal documents with contextual chunking"""
|
||||
start_time = time.time()
|
||||
|
||||
# Clean up first
|
||||
self.cleanup_existing_index()
|
||||
|
||||
# Get all documents
|
||||
all_documents = self.get_all_legal_documents()
|
||||
|
||||
# Filter by categories if specified
|
||||
if categories:
|
||||
all_documents = [d for d in all_documents if any(cat in d["category"] for cat in categories)]
|
||||
|
||||
# Limit documents if specified
|
||||
if max_docs:
|
||||
all_documents = all_documents[:max_docs]
|
||||
|
||||
logger.info(f"Processing {len(all_documents)} documents...")
|
||||
|
||||
for i, doc_info in enumerate(all_documents):
|
||||
logger.info(f"\n[{i+1}/{len(all_documents)}] Processing: {doc_info['name']}")
|
||||
logger.info(f" Category: {doc_info['category']}")
|
||||
|
||||
# Track category
|
||||
self.stats["categories_processed"].add(doc_info["category"])
|
||||
|
||||
# Read document
|
||||
content = self.read_document(doc_info)
|
||||
if not content:
|
||||
continue
|
||||
|
||||
# Process with contextual chunking and immediate indexing
|
||||
try:
|
||||
# Process and index chunks immediately as they're generated
|
||||
chunks = self.process_document(doc_info, content, index_immediately=True)
|
||||
self.stats["chunks_created"] += len(chunks)
|
||||
|
||||
# Log contextual stats
|
||||
if self.use_contextual:
|
||||
contextual_count = sum(1 for c in chunks if c.context)
|
||||
logger.info(f" ✓ Created and indexed {len(chunks)} chunks ({contextual_count} with context)")
|
||||
else:
|
||||
logger.info(f" ✓ Created and indexed {len(chunks)} chunks (no context)")
|
||||
|
||||
# Save document info immediately after processing
|
||||
self.save_document_info(doc_info, chunks)
|
||||
|
||||
self.stats["documents_processed"] += 1
|
||||
|
||||
# Update token statistics
|
||||
if self.use_contextual:
|
||||
self.stats["total_context_tokens"] += sum(c.context_tokens for c in chunks)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing {doc_info['name']}: {e}")
|
||||
self.stats["errors"] += 1
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Get final chunker statistics
|
||||
chunker_stats = self.chunker.get_statistics()
|
||||
|
||||
# Print statistics
|
||||
self._print_statistics(elapsed, chunker_stats)
|
||||
|
||||
def _print_statistics(self, elapsed_time: float, chunker_stats: Dict[str, Any]):
|
||||
"""Print processing statistics"""
|
||||
print("\n" + "="*60)
|
||||
print("CONTEXTUAL INDEXING COMPLETE")
|
||||
print("="*60)
|
||||
print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
|
||||
print(f"Time elapsed: {elapsed_time:.2f} seconds")
|
||||
print(f"Categories processed: {len(self.stats['categories_processed'])}")
|
||||
|
||||
# Show categories
|
||||
categories_list = sorted(self.stats['categories_processed'])
|
||||
for cat in categories_list:
|
||||
print(f" • {cat}")
|
||||
|
||||
print(f"\nDocuments processed: {self.stats['documents_processed']}")
|
||||
print(f"Chunks created: {self.stats['chunks_created']}")
|
||||
|
||||
if self.use_contextual:
|
||||
print(f"Contextual chunks: {self.stats['contextual_chunks']}")
|
||||
print(f"Chunks indexed: {self.stats['chunks_indexed']}")
|
||||
|
||||
# Token usage
|
||||
print(f"\nContext Generation Statistics:")
|
||||
print(f" Total context tokens: {chunker_stats.get('total_context_tokens', 0):,}")
|
||||
print(f" Average tokens per chunk: {chunker_stats.get('avg_context_tokens', 0):.1f}")
|
||||
print(f" Total generation time: {chunker_stats.get('total_generation_time', 0):.2f}s")
|
||||
print(f" Average time per chunk: {chunker_stats.get('avg_generation_time', 0):.2f}s")
|
||||
|
||||
# Cache statistics
|
||||
print(f"\nCache Performance:")
|
||||
print(f" Cache hits: {chunker_stats.get('cache_hits', 0)}")
|
||||
print(f" Cache misses: {chunker_stats.get('cache_misses', 0)}")
|
||||
print(f" Hit rate: {chunker_stats.get('cache_hit_rate', 0):.1%}")
|
||||
|
||||
# Cost estimation
|
||||
print(f"\nCost Estimation:")
|
||||
print(f" Estimated cost: ${chunker_stats.get('estimated_cost', 0):.2f}")
|
||||
else:
|
||||
print(f"Chunks indexed: {self.stats['chunks_indexed']}")
|
||||
|
||||
print(f"\nErrors: {self.stats['errors']}")
|
||||
|
||||
if self.stats['chunks_created'] > 0:
|
||||
avg_chunks = self.stats['chunks_created'] / max(1, self.stats['documents_processed'])
|
||||
print(f"Average chunks per document: {avg_chunks:.1f}")
|
||||
|
||||
if elapsed_time > 0 and self.stats['documents_processed'] > 0:
|
||||
docs_per_min = (self.stats['documents_processed'] / elapsed_time) * 60
|
||||
print(f"Processing speed: {docs_per_min:.1f} docs/minute")
|
||||
|
||||
print("="*60 + "\n")
|
||||
|
||||
def compare_retrieval_methods(self, test_queries: Optional[List[str]] = None):
|
||||
"""Compare contextual vs non-contextual retrieval"""
|
||||
if not test_queries:
|
||||
test_queries = [
|
||||
"什么是合同的成立条件",
|
||||
"劳动者的基本权利有哪些",
|
||||
"如何处理交通事故责任",
|
||||
"什么是正当防卫",
|
||||
"公司股东的权利和义务"
|
||||
]
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("RETRIEVAL COMPARISON TEST")
|
||||
print("="*60)
|
||||
print(f"Mode: {'Contextual' if self.use_contextual else 'Non-contextual'}")
|
||||
|
||||
for query in test_queries:
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.pipeline_url}/search",
|
||||
json={
|
||||
"query": query,
|
||||
"mode": "hybrid",
|
||||
"top_k": 10,
|
||||
"rerank_top_k": 5
|
||||
}, timeout=30
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
results = response.json()
|
||||
print(f"\n✓ Query: '{query}'")
|
||||
|
||||
if "results" in results:
|
||||
print(f" Found {len(results['results'])} results")
|
||||
|
||||
# Show top 3 results
|
||||
for i, result in enumerate(results['results'][:3]):
|
||||
score = result.get('score', result.get('rerank_score', 'N/A'))
|
||||
metadata = result.get('metadata', {})
|
||||
|
||||
print(f"\n Result {i+1}:")
|
||||
print(f" Score: {score}")
|
||||
print(f" Category: {metadata.get('category', 'Unknown')}")
|
||||
print(f" Document: {metadata.get('doc_title', 'Unknown')}")
|
||||
|
||||
# Show context if available
|
||||
if metadata.get('context'):
|
||||
print(f" Context: {metadata['context'][:100]}...")
|
||||
|
||||
# Show text preview
|
||||
text = result.get('text', '')
|
||||
# If contextual, try to extract original text
|
||||
if metadata.get('contextual') and metadata.get('original_text'):
|
||||
text = metadata['original_text']
|
||||
|
||||
print(f" Preview: {text[:150]}...")
|
||||
else:
|
||||
print(f" No results found")
|
||||
else:
|
||||
print(f"✗ Query '{query}' failed: {response.status_code}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"✗ Error testing '{query}': {e}")
|
||||
|
||||
print("\n" + "="*60 + "\n")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Index local legal documents with contextual retrieval")
|
||||
parser.add_argument("--pipeline-url", default=RETRIEVAL_PIPELINE_URL, help="Retrieval pipeline URL")
|
||||
parser.add_argument("--max-docs", type=int, help="Maximum number of documents to process")
|
||||
parser.add_argument("--categories", nargs="+", help="Specific categories to process")
|
||||
parser.add_argument("--no-contextual", action="store_true", help="Disable contextual enhancement")
|
||||
parser.add_argument("--no-cleanup", action="store_true", help="Don't clean existing indexes")
|
||||
parser.add_argument("--compare", action="store_true", help="Run comparison test after indexing")
|
||||
parser.add_argument("--batch-size", type=int, default=10, help="Batch size for indexing")
|
||||
parser.add_argument("--llm-provider", default="kimi", help="LLM provider for context generation")
|
||||
parser.add_argument("--llm-model", help="Specific LLM model to use")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Configure LLM if specified
|
||||
llm_config = None
|
||||
if args.llm_provider or args.llm_model:
|
||||
llm_config = LLMConfig(
|
||||
provider=args.llm_provider,
|
||||
model=args.llm_model
|
||||
)
|
||||
|
||||
# Create indexer
|
||||
indexer = ContextualLegalIndexer(
|
||||
pipeline_url=args.pipeline_url,
|
||||
use_contextual=not args.no_contextual,
|
||||
llm_config=llm_config
|
||||
)
|
||||
|
||||
# Skip cleanup if requested
|
||||
if args.no_cleanup:
|
||||
indexer.cleanup_existing_index = lambda: logger.info("Skipping cleanup (--no-cleanup flag)")
|
||||
|
||||
# Process documents
|
||||
indexer.process_all_documents(
|
||||
max_docs=args.max_docs,
|
||||
categories=args.categories,
|
||||
batch_size=args.batch_size
|
||||
)
|
||||
|
||||
# Run comparison test if requested
|
||||
if args.compare:
|
||||
indexer.compare_retrieval_methods()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+1
@@ -0,0 +1 @@
|
||||
../agentic-rag/laws
|
||||
@@ -0,0 +1,292 @@
|
||||
"""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", timeout=30)
|
||||
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 Exception:
|
||||
logger.warning("Cannot connect to local retrieval pipeline")
|
||||
logger.info("Will continue anyway - searches may fail")
|
||||
|
||||
elif config.knowledge_base.type == KnowledgeBaseType.DIFY:
|
||||
if not config.knowledge_base.dify_api_key:
|
||||
logger.warning("Dify API key not set")
|
||||
logger.info("Please set DIFY_API_KEY environment variable")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def run_interactive_mode(agent: AgenticRAG, mode: str = "agentic"):
|
||||
"""Run interactive query mode"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Agentic RAG System - {mode.capitalize()} Mode")
|
||||
print(f"Verbose: {'Enabled' if agent.config.agent.verbose else 'Disabled'} | Top-K: {agent.config.knowledge_base.local_top_k}")
|
||||
print(f"{'='*60}")
|
||||
print("Type 'quit' or 'exit' to stop")
|
||||
print("Type 'clear' to clear conversation history")
|
||||
print("Type 'mode' to switch between agentic/non-agentic modes")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
current_mode = mode
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input("\n[USER] > ").strip()
|
||||
|
||||
if user_input.lower() in ['quit', 'exit']:
|
||||
print("\nGoodbye!")
|
||||
break
|
||||
|
||||
if user_input.lower() == 'clear':
|
||||
agent.clear_history()
|
||||
print("Conversation history cleared.")
|
||||
continue
|
||||
|
||||
if user_input.lower() == 'mode':
|
||||
current_mode = "non-agentic" if current_mode == "agentic" else "agentic"
|
||||
print(f"Switched to {current_mode} mode")
|
||||
continue
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
# Process query
|
||||
print(f"\n[ASSISTANT ({current_mode})] > ", end="", flush=True)
|
||||
|
||||
if current_mode == "agentic":
|
||||
response = agent.query(user_input, stream=True)
|
||||
else:
|
||||
response = agent.query_non_agentic(user_input, stream=True)
|
||||
|
||||
# Handle streaming response
|
||||
if hasattr(response, '__iter__'):
|
||||
for chunk in response:
|
||||
print(chunk, end="", flush=True)
|
||||
print() # New line after response
|
||||
else:
|
||||
print(response)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\nInterrupted. Type 'quit' to exit.")
|
||||
except Exception as e:
|
||||
logger.error(f"Error: {e}")
|
||||
print(f"\nError processing query: {e}")
|
||||
|
||||
|
||||
def run_batch_mode(agent: AgenticRAG, queries_file: str, output_file: str, mode: str = "agentic"):
|
||||
"""Run batch queries from file"""
|
||||
try:
|
||||
with open(queries_file, 'r', encoding='utf-8') as f:
|
||||
queries = [line.strip() for line in f if line.strip()]
|
||||
except Exception as e:
|
||||
logger.error(f"Error reading queries file: {e}")
|
||||
return
|
||||
|
||||
results = []
|
||||
|
||||
for i, query in enumerate(queries, 1):
|
||||
print(f"\n[{i}/{len(queries)}] Processing: {query[:100]}...")
|
||||
|
||||
try:
|
||||
if mode == "agentic":
|
||||
response = agent.query(query, stream=False)
|
||||
else:
|
||||
response = agent.query_non_agentic(query, stream=False)
|
||||
|
||||
results.append({
|
||||
"query": query,
|
||||
"response": response,
|
||||
"mode": mode
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing query: {e}")
|
||||
results.append({
|
||||
"query": query,
|
||||
"response": f"Error: {str(e)}",
|
||||
"mode": mode
|
||||
})
|
||||
|
||||
# Save results
|
||||
try:
|
||||
with open(output_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
print(f"\nResults saved to {output_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving results: {e}")
|
||||
|
||||
|
||||
def run_comparison_mode(agent: AgenticRAG, query: str):
|
||||
"""Run both modes and compare results"""
|
||||
print(f"\n{'='*60}")
|
||||
print("Comparison Mode - Running both Agentic and Non-Agentic")
|
||||
print(f"{'='*60}")
|
||||
print(f"Query: {query}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
# Run non-agentic mode
|
||||
print("\n[NON-AGENTIC MODE]")
|
||||
print("-" * 40)
|
||||
non_agentic_response = agent.query_non_agentic(query, stream=False)
|
||||
print(non_agentic_response)
|
||||
|
||||
# Clear history for fair comparison
|
||||
agent.clear_history()
|
||||
|
||||
# Run agentic mode
|
||||
print("\n[AGENTIC MODE]")
|
||||
print("-" * 40)
|
||||
agentic_response = agent.query(query, stream=False)
|
||||
print(agentic_response)
|
||||
|
||||
print(f"\n{'='*60}")
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function"""
|
||||
parser = argparse.ArgumentParser(description="Agentic RAG System")
|
||||
|
||||
# Mode selection
|
||||
parser.add_argument("--mode", choices=["agentic", "non-agentic", "compare"],
|
||||
default="agentic", help="Query mode")
|
||||
|
||||
# Query options
|
||||
parser.add_argument("--query", type=str, help="Single query to process")
|
||||
parser.add_argument("--batch", type=str, help="Path to file with queries (one per line)")
|
||||
parser.add_argument("--output", type=str, default="results.json",
|
||||
help="Output file for batch results")
|
||||
|
||||
# Configuration options
|
||||
parser.add_argument("--provider", type=str, help="LLM provider")
|
||||
parser.add_argument("--model", type=str, help="LLM model")
|
||||
parser.add_argument("--kb-type", choices=["local", "dify"], help="Knowledge base type")
|
||||
parser.add_argument("--verbose", action="store_true", help="Verbose output (default: True)")
|
||||
parser.add_argument("--no-verbose", action="store_true", help="Disable verbose output")
|
||||
|
||||
# Indexing options
|
||||
parser.add_argument("--index", type=str, help="Path to index (file or directory)")
|
||||
parser.add_argument("--chunk-size", type=int, default=2048, help="Chunk size for indexing")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Setup environment
|
||||
if not setup_environment():
|
||||
logger.warning("Environment setup incomplete, continuing anyway...")
|
||||
|
||||
# Load or create config
|
||||
config = Config.from_env()
|
||||
|
||||
# Set verbose mode by default (can be disabled with --no-verbose)
|
||||
config.agent.verbose = True # Default to verbose mode
|
||||
|
||||
# Override config with command line args
|
||||
if args.provider:
|
||||
config.llm.provider = args.provider
|
||||
if args.model:
|
||||
config.llm.model = args.model
|
||||
if args.kb_type:
|
||||
config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
|
||||
|
||||
# Handle verbose mode (default is True, can be disabled with --no-verbose)
|
||||
if args.no_verbose:
|
||||
config.agent.verbose = False
|
||||
elif args.verbose:
|
||||
config.agent.verbose = True # Explicitly set if --verbose is used
|
||||
|
||||
# Handle indexing if requested
|
||||
if args.index:
|
||||
print(f"\n{'='*60}")
|
||||
print("Indexing Documents")
|
||||
print(f"{'='*60}")
|
||||
|
||||
config.chunking.chunk_size = args.chunk_size
|
||||
indexer = DocumentIndexer(config.knowledge_base, config.chunking)
|
||||
|
||||
from pathlib import Path
|
||||
path = Path(args.index)
|
||||
|
||||
if path.is_file():
|
||||
result = indexer.index_file(str(path))
|
||||
elif path.is_dir():
|
||||
result = indexer.index_directory(str(path))
|
||||
else:
|
||||
print(f"Path not found: {path}")
|
||||
return
|
||||
|
||||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
# Create agent
|
||||
agent = AgenticRAG(config)
|
||||
|
||||
# Handle different execution modes
|
||||
if args.query and args.mode == "compare":
|
||||
# Comparison mode with single query
|
||||
run_comparison_mode(agent, args.query)
|
||||
|
||||
elif args.query:
|
||||
# Single query mode
|
||||
print(f"\n[Query] {args.query}")
|
||||
print(f"[Mode] {args.mode}")
|
||||
print(f"[Verbose] {'Enabled' if config.agent.verbose else 'Disabled'}")
|
||||
print(f"[Top-K] {config.knowledge_base.local_top_k}")
|
||||
print("-" * 40)
|
||||
|
||||
if args.mode == "agentic":
|
||||
response = agent.query(args.query, stream=False)
|
||||
else:
|
||||
response = agent.query_non_agentic(args.query, stream=False)
|
||||
|
||||
print(response)
|
||||
|
||||
elif args.batch:
|
||||
# Batch mode
|
||||
run_batch_mode(agent, args.batch, args.output, args.mode)
|
||||
|
||||
else:
|
||||
# Interactive mode (default)
|
||||
run_interactive_mode(agent, args.mode)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Quick start script to test the Contextual Retrieval System
|
||||
|
||||
This script provides a quick way to test contextual retrieval
|
||||
with a sample document and see the improvements.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from config import Config
|
||||
from contextual_chunking import ContextualChunker
|
||||
from contextual_tools import ContextualKnowledgeBaseTools
|
||||
|
||||
# Simple logging for quickstart
|
||||
logging.basicConfig(level=logging.INFO, format='%(message)s')
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main():
|
||||
"""Quick demonstration of contextual retrieval"""
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("CONTEXTUAL RETRIEVAL - QUICK START")
|
||||
print("="*60 + "\n")
|
||||
|
||||
# Sample document about multiple companies
|
||||
document = """
|
||||
2023 Technology Sector Report
|
||||
|
||||
Apple Inc. Performance:
|
||||
Apple reported exceptional results in 2023. The company's revenue reached
|
||||
$394 billion, with iPhone sales contributing 52% of total revenue. The
|
||||
services division showed strong growth of 16% year-over-year. Tim Cook
|
||||
emphasized the company's commitment to innovation and sustainability.
|
||||
|
||||
Microsoft Corporation Update:
|
||||
Microsoft achieved record cloud revenue in 2023. Azure revenue grew by 27%
|
||||
as enterprises accelerated digital transformation. The company's total
|
||||
revenue was $211 billion. CEO Satya Nadella highlighted AI integration
|
||||
across all product lines as a key strategic priority.
|
||||
|
||||
Google (Alphabet) Highlights:
|
||||
Google's parent company Alphabet reported $283 billion in revenue for 2023.
|
||||
Search advertising remained the largest revenue driver at $175 billion.
|
||||
The company increased AI research spending by 30% to maintain competitive
|
||||
advantage. YouTube advertising revenue exceeded $40 billion.
|
||||
|
||||
Market Analysis:
|
||||
The technology sector showed resilience despite economic headwinds. Companies
|
||||
that invested heavily in AI and cloud infrastructure outperformed the market.
|
||||
The sector's average growth rate was 12%, with cloud services growing at 25%
|
||||
and traditional hardware declining by 3%.
|
||||
"""
|
||||
|
||||
print("Step 1: Initializing systems...")
|
||||
config = Config.from_env()
|
||||
|
||||
# Create both contextual and non-contextual systems
|
||||
contextual_chunker = ContextualChunker(use_contextual=True)
|
||||
non_contextual_chunker = ContextualChunker(use_contextual=False)
|
||||
|
||||
contextual_kb = ContextualKnowledgeBaseTools(use_contextual=True)
|
||||
non_contextual_kb = ContextualKnowledgeBaseTools(use_contextual=False)
|
||||
|
||||
print("\nStep 2: Processing document...")
|
||||
print("-" * 40)
|
||||
|
||||
# Process with contextual system
|
||||
print("Creating contextual chunks...")
|
||||
contextual_chunks = contextual_chunker.chunk_document(
|
||||
text=document,
|
||||
doc_id="tech_report_2023"
|
||||
)
|
||||
contextual_kb.index_contextual_chunks(contextual_chunks)
|
||||
print(f"✓ Created {len(contextual_chunks)} contextual chunks")
|
||||
|
||||
# Process with non-contextual system
|
||||
print("Creating non-contextual chunks...")
|
||||
non_contextual_chunks = non_contextual_chunker.chunk_document(
|
||||
text=document,
|
||||
doc_id="tech_report_2023"
|
||||
)
|
||||
non_contextual_kb.index_contextual_chunks(non_contextual_chunks)
|
||||
print(f"✓ Created {len(non_contextual_chunks)} non-contextual chunks")
|
||||
|
||||
# Show example contextual chunk
|
||||
if contextual_chunks:
|
||||
print("\nExample Contextual Chunk:")
|
||||
print("-" * 40)
|
||||
chunk = contextual_chunks[0]
|
||||
print(f"Original text: {chunk.text[:100]}...")
|
||||
print(f"Added context: {chunk.context}")
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("Step 3: Testing Search Queries")
|
||||
print("="*60)
|
||||
|
||||
# Test queries
|
||||
queries = [
|
||||
"What was the company's revenue?",
|
||||
"Which company emphasized AI?",
|
||||
"What was the growth rate?"
|
||||
]
|
||||
|
||||
for query in queries:
|
||||
print(f"\nQuery: '{query}'")
|
||||
print("-" * 40)
|
||||
|
||||
# Contextual search
|
||||
contextual_results = contextual_kb.contextual_search(query, top_k=1)
|
||||
|
||||
# Non-contextual search
|
||||
non_contextual_results = non_contextual_kb.contextual_search(query, top_k=1)
|
||||
|
||||
print("\nContextual Result:")
|
||||
if contextual_results:
|
||||
result = contextual_results[0]
|
||||
print(f" Score: {result.score:.4f}")
|
||||
if result.context_text:
|
||||
print(f" Context: {result.context_text[:80]}...")
|
||||
print(f" Match: {result.text[:100]}...")
|
||||
else:
|
||||
print(" No results")
|
||||
|
||||
print("\nNon-Contextual Result:")
|
||||
if non_contextual_results:
|
||||
result = non_contextual_results[0]
|
||||
print(f" Score: {result.score:.4f}")
|
||||
print(f" Match: {result.text[:100]}...")
|
||||
else:
|
||||
print(" No results")
|
||||
|
||||
# Compare scores
|
||||
if contextual_results and non_contextual_results:
|
||||
improvement = ((contextual_results[0].score - non_contextual_results[0].score)
|
||||
/ non_contextual_results[0].score * 100)
|
||||
print(f"\n📊 Improvement: {improvement:+.1f}%")
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("SUMMARY")
|
||||
print("="*60)
|
||||
|
||||
# Get statistics
|
||||
stats = contextual_chunker.get_statistics()
|
||||
|
||||
print(f"\nContextual Chunking Statistics:")
|
||||
print(f" Chunks processed: {stats['total_chunks']}")
|
||||
print(f" Context tokens used: {stats['total_context_tokens']}")
|
||||
print(f" Estimated cost: ${stats['estimated_cost']:.4f}")
|
||||
|
||||
print("\nKey Insights:")
|
||||
print("✓ Contextual chunks preserve company-specific information")
|
||||
print("✓ Ambiguous queries ('the company') are resolved correctly")
|
||||
print("✓ Search accuracy improves significantly with context")
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("Quick start complete! Try with your own documents:")
|
||||
print(" python contextual_main.py --mode index --document your_file.txt")
|
||||
print("="*60 + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,8 @@
|
||||
openai>=1.0.0
|
||||
requests>=2.31.0
|
||||
python-dotenv>=1.0.0
|
||||
pydantic>=2.0.0
|
||||
tiktoken>=0.5.0
|
||||
numpy>=1.24.0
|
||||
rank-bm25>=0.2.2
|
||||
scikit-learn>=1.3.0
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
"""conversation_history_limit=0 must omit history, not include all via [-0:]."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from agent import AgenticRAG
|
||||
|
||||
|
||||
def test_history_limit_zero_omits_history():
|
||||
agent = object.__new__(AgenticRAG)
|
||||
agent.config = SimpleNamespace(
|
||||
agent=SimpleNamespace(conversation_history_limit=0)
|
||||
)
|
||||
agent._get_system_prompt = lambda: "sys"
|
||||
agent.conversation_history = [
|
||||
{"role": "user", "content": "old1"},
|
||||
{"role": "assistant", "content": "old2"},
|
||||
]
|
||||
msgs = agent._build_messages("new question")
|
||||
assert msgs == [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "new question"},
|
||||
]
|
||||
|
||||
|
||||
def test_positive_history_limit_still_keeps_tail():
|
||||
agent = object.__new__(AgenticRAG)
|
||||
agent.config = SimpleNamespace(
|
||||
agent=SimpleNamespace(conversation_history_limit=1)
|
||||
)
|
||||
agent._get_system_prompt = lambda: "sys"
|
||||
agent.conversation_history = [
|
||||
{"role": "user", "content": "old1"},
|
||||
{"role": "assistant", "content": "old2"},
|
||||
]
|
||||
msgs = agent._build_messages("new question")
|
||||
assert msgs == [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "assistant", "content": "old2"},
|
||||
{"role": "user", "content": "new question"},
|
||||
]
|
||||
@@ -0,0 +1,195 @@
|
||||
#!/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 _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_basic_functionality():
|
||||
"""Pytest contract: failures are assertions, not a returned status value."""
|
||||
assert _basic_functionality()
|
||||
|
||||
|
||||
def _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
|
||||
|
||||
|
||||
def test_evaluation_dataset():
|
||||
"""Pytest contract: failures are assertions, not a returned status value."""
|
||||
assert _evaluation_dataset()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("🚀 Agentic RAG System - Test Suite")
|
||||
print("="*60)
|
||||
|
||||
# Run tests
|
||||
success = _basic_functionality()
|
||||
|
||||
if success:
|
||||
_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)")
|
||||
@@ -0,0 +1,518 @@
|
||||
"""Tools for knowledge base interaction"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import requests
|
||||
from typing import Dict, Any, List, Optional
|
||||
from dataclasses import dataclass
|
||||
from config import KnowledgeBaseConfig, KnowledgeBaseType
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchResult:
|
||||
"""Search result from knowledge base"""
|
||||
doc_id: str
|
||||
chunk_id: str
|
||||
text: str
|
||||
score: float
|
||||
metadata: Dict[str, Any] = None
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"doc_id": self.doc_id,
|
||||
"chunk_id": self.chunk_id,
|
||||
"text": self.text,
|
||||
"score": self.score,
|
||||
"metadata": self.metadata or {}
|
||||
}
|
||||
|
||||
|
||||
class KnowledgeBaseTools:
|
||||
"""Tools for interacting with knowledge base"""
|
||||
|
||||
def __init__(self, config: KnowledgeBaseConfig):
|
||||
self.config = config
|
||||
self.document_store = {} # In-memory store for documents
|
||||
|
||||
# Load document store if exists
|
||||
try:
|
||||
with open(config.document_store_path, 'r', encoding='utf-8') as f:
|
||||
self.document_store = json.load(f)
|
||||
except FileNotFoundError:
|
||||
logger.info("No existing document store found, starting fresh")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading document store: {e}")
|
||||
|
||||
def save_document_store(self):
|
||||
"""Save document store to disk"""
|
||||
try:
|
||||
with open(self.config.document_store_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.document_store, f, ensure_ascii=False, indent=2)
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving document store: {e}")
|
||||
|
||||
def knowledge_base_search(self, query: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search the knowledge base with a natural language query.
|
||||
|
||||
Args:
|
||||
query: Natural language query string
|
||||
|
||||
Returns:
|
||||
List of matching document chunks with scores
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Knowledge base search initiated - Type: {self.config.type}, Query: '{query}'")
|
||||
|
||||
if self.config.type == KnowledgeBaseType.LOCAL:
|
||||
return self._search_local(query)
|
||||
elif self.config.type == KnowledgeBaseType.DIFY:
|
||||
return self._search_dify(query)
|
||||
elif self.config.type == KnowledgeBaseType.RAPTOR:
|
||||
return self._search_raptor(query)
|
||||
elif self.config.type == KnowledgeBaseType.GRAPHRAG:
|
||||
return self._search_graphrag(query)
|
||||
else:
|
||||
logger.error(f"Unsupported knowledge base type: {self.config.type}")
|
||||
raise ValueError(f"Unsupported knowledge base type: {self.config.type}")
|
||||
except ValueError:
|
||||
raise # Re-raise ValueError for unsupported types
|
||||
except Exception as e:
|
||||
logger.error(f"Error in knowledge base search: {e}")
|
||||
return []
|
||||
|
||||
def _search_local(self, query: str) -> List[Dict[str, Any]]:
|
||||
"""Search using local retrieval pipeline"""
|
||||
try:
|
||||
logger.info(f"Searching local knowledge base for: {query}")
|
||||
|
||||
response = requests.post(
|
||||
f"{self.config.local_base_url}/search",
|
||||
json={
|
||||
"query": query,
|
||||
"mode": "hybrid",
|
||||
"top_k": self.config.local_top_k,
|
||||
"rerank": True
|
||||
},
|
||||
timeout=30 # Add timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
results = []
|
||||
data = response.json()
|
||||
|
||||
# The retrieval pipeline returns results in 'reranked_results' field for hybrid mode
|
||||
results_field = data.get("reranked_results") or data.get("results") or []
|
||||
|
||||
if not results_field:
|
||||
# Check other possible fields
|
||||
if data.get("dense_results"):
|
||||
results_field = data.get("dense_results", [])
|
||||
elif data.get("sparse_results"):
|
||||
results_field = data.get("sparse_results", [])
|
||||
|
||||
if not results_field:
|
||||
logger.warning(f"Search returned empty results for query: {query}")
|
||||
logger.debug(f"Response keys: {data.keys()}")
|
||||
return []
|
||||
|
||||
for item in results_field:
|
||||
# Extract doc_id - the field name varies between results
|
||||
doc_id = item.get("doc_id", "")
|
||||
|
||||
# Use doc_id as chunk_id since the retrieval pipeline indexes chunks
|
||||
chunk_id = doc_id
|
||||
|
||||
# Get text from the result
|
||||
text = item.get("text", "")
|
||||
|
||||
# Get score - might be 'rerank_score' or 'score'
|
||||
score = item.get("rerank_score", item.get("score", 0.0))
|
||||
|
||||
result = SearchResult(
|
||||
doc_id=doc_id,
|
||||
chunk_id=chunk_id,
|
||||
text=text,
|
||||
score=score,
|
||||
metadata=item.get("metadata", {})
|
||||
)
|
||||
results.append(result.to_dict())
|
||||
|
||||
logger.info(f"Local search returned {len(results)} results for query: {query}")
|
||||
return results
|
||||
|
||||
except requests.exceptions.Timeout as e:
|
||||
logger.error(f"Timeout connecting to local retrieval pipeline at {self.config.local_base_url}: {e}")
|
||||
return []
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
logger.error(f"Cannot connect to local retrieval pipeline at {self.config.local_base_url}: {e}")
|
||||
logger.info("Make sure the retrieval pipeline is running: cd ../retrieval-pipeline && python main.py")
|
||||
return []
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Error connecting to local retrieval pipeline: {e}")
|
||||
return []
|
||||
|
||||
def _search_dify(self, query: str) -> List[Dict[str, Any]]:
|
||||
"""Search using Dify API"""
|
||||
if not self.config.dify_api_key:
|
||||
logger.error("Dify API key not configured")
|
||||
return []
|
||||
|
||||
try:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.config.dify_api_key}",
|
||||
"Content-Type": "application/json"
|
||||
}
|
||||
|
||||
payload = {
|
||||
"query": query,
|
||||
"top_k": self.config.dify_top_k
|
||||
}
|
||||
|
||||
if self.config.dify_dataset_id:
|
||||
payload["dataset_id"] = self.config.dify_dataset_id
|
||||
|
||||
response = requests.post(
|
||||
f"{self.config.dify_base_url}/datasets/search",
|
||||
headers=headers,
|
||||
json=payload, timeout=30
|
||||
)
|
||||
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
|
||||
}, timeout=30
|
||||
)
|
||||
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
|
||||
}, timeout=30
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
results = []
|
||||
data = response.json()
|
||||
|
||||
for i, item in enumerate(data.get("results", [])):
|
||||
# GraphRAG returns entities or communities
|
||||
result_type = item.get("type", "unknown")
|
||||
|
||||
if result_type == "entity":
|
||||
doc_id = item.get("id", f"entity_{i}")
|
||||
chunk_id = f"{doc_id}_{item.get('entity_type', 'unknown')}"
|
||||
text_content = f"{item.get('name', '')}. {item.get('description', '')}"
|
||||
metadata = {
|
||||
"type": "entity",
|
||||
"entity_type": item.get("entity_type"),
|
||||
"related_entities": item.get("related_entities", [])
|
||||
}
|
||||
else: # community
|
||||
doc_id = item.get("id", f"community_{i}")
|
||||
chunk_id = f"{doc_id}_level_{item.get('level', 0)}"
|
||||
text_content = item.get("summary", "")
|
||||
metadata = {
|
||||
"type": "community",
|
||||
"level": item.get("level", 0),
|
||||
"entity_count": item.get("entity_count", 0),
|
||||
"sample_entities": item.get("sample_entities", [])
|
||||
}
|
||||
|
||||
result = SearchResult(
|
||||
doc_id=doc_id,
|
||||
chunk_id=chunk_id,
|
||||
text=text_content,
|
||||
score=item.get("score", 0.0),
|
||||
metadata={**metadata, "source": "graphrag"}
|
||||
)
|
||||
results.append(result.to_dict())
|
||||
|
||||
logger.info(f"GraphRAG search returned {len(results)} results")
|
||||
return results
|
||||
|
||||
except requests.exceptions.RequestException as e:
|
||||
logger.error(f"Error connecting to GraphRAG index: {e}")
|
||||
return []
|
||||
|
||||
def get_document(self, doc_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Retrieve the entire document from the knowledge base.
|
||||
|
||||
Args:
|
||||
doc_id: Document ID
|
||||
|
||||
Returns:
|
||||
Full document content and metadata
|
||||
"""
|
||||
try:
|
||||
# First check local document store
|
||||
if doc_id in self.document_store:
|
||||
return self.document_store[doc_id]
|
||||
|
||||
if self.config.type == KnowledgeBaseType.LOCAL:
|
||||
return self._get_document_local(doc_id)
|
||||
elif self.config.type == KnowledgeBaseType.DIFY:
|
||||
return self._get_document_dify(doc_id)
|
||||
elif self.config.type == KnowledgeBaseType.RAPTOR:
|
||||
return self._get_document_raptor(doc_id)
|
||||
elif self.config.type == KnowledgeBaseType.GRAPHRAG:
|
||||
return self._get_document_graphrag(doc_id)
|
||||
else:
|
||||
raise ValueError(f"Unsupported knowledge base type: {self.config.type}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error retrieving document {doc_id}: {e}")
|
||||
return {"error": f"Document {doc_id} not found"}
|
||||
|
||||
def _get_document_local(self, doc_id: str) -> Dict[str, Any]:
|
||||
"""Get document from local retrieval pipeline"""
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{self.config.local_base_url}/documents/{doc_id}", timeout=30
|
||||
)
|
||||
|
||||
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, timeout=30
|
||||
)
|
||||
|
||||
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
|
||||
}, timeout=30
|
||||
)
|
||||
|
||||
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"
|
||||
}, timeout=30
|
||||
)
|
||||
|
||||
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"]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,130 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-10",
|
||||
"run_id": "20260729T205807Z-3_10-35531780",
|
||||
"created_at": "2026-07-29T20:58:07.546305+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/validation/runs/20260729T205807Z-3_10-35531780",
|
||||
"artifacts": {
|
||||
"evidence.json": "e8d2ea403646d0ed0041a5849f38bf16d43481ea8b622a30ae9bdf774b60b11b",
|
||||
"receipts.json": "8cf26248f9fd99206784831992e03d67895e51d92f3a2d4223ed56e22b8a6ca2",
|
||||
"manifest.json": "fa3b31815a1155db02c1ad4057bccbbe1bd83fb2f14310b5963c0d7949ccb61c"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/campaign.py",
|
||||
"sha256": "2e61651942f869a2374f24fe40e92aaf332a5bd5abf9c3dbc951fff3a31d9e28",
|
||||
"bytes": 13270
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/compare_retrieval.py",
|
||||
"sha256": "5cb1213c44c94a875e14475c3f371799dd750c00001a9ab36b49034b460f94ef",
|
||||
"bytes": 15372
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/document_store.json",
|
||||
"sha256": "1fe9c0a7096fc9fcf50c2c5c8ea2008b001243932043ab8882cff436671f8d48",
|
||||
"bytes": 154529
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/evaluation/retrieval_eval.json",
|
||||
"sha256": "7df53f5877ac6a8ec96aa3b2dc575e8cc556e1c9a3eb0ce28800ddd20d0f4a99",
|
||||
"bytes": 4426
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/1-宪法/宪法.md",
|
||||
"sha256": "4cbc5d5d8a4632eb5587ab6377ae5251ce21a149716690864db0b0ab6cd4b386",
|
||||
"bytes": 52540
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/2-宪法相关法/检察官法(2019-04-23).md",
|
||||
"sha256": "16285739d98e9661208bf8b1f334d40d902af794dcef322605bc93ffbfb7c530",
|
||||
"bytes": 20824
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"methods": {
|
||||
"plain_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.6,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.7508333333333334
|
||||
},
|
||||
"contextual_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8555555555555555
|
||||
},
|
||||
"plain_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9222222222222223
|
||||
},
|
||||
"contextual_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.9333333333333333,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9666666666666667
|
||||
},
|
||||
"plain_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.7333333333333333,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8444444444444444
|
||||
},
|
||||
"contextual_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9133333333333333
|
||||
}
|
||||
},
|
||||
"index_time": {
|
||||
"context_generation_ms": 61804.353,
|
||||
"embedding_ms": 365453.495,
|
||||
"usage": {
|
||||
"prompt_tokens": 219888,
|
||||
"completion_tokens": 8846,
|
||||
"total_tokens": 228734
|
||||
},
|
||||
"estimated_cost_usd": 0.033918,
|
||||
"pricing_assumption": {
|
||||
"input_per_million_usd": 0.11,
|
||||
"output_per_million_usd": 1.1
|
||||
}
|
||||
},
|
||||
"errors": 0
|
||||
},
|
||||
"acceptance": {
|
||||
"live_prefix_for_every_chunk": true,
|
||||
"full_source_document_and_target_chunk_in_requests": true,
|
||||
"same_chunks_and_queries": true,
|
||||
"plain_contextual_bm25_dense_hybrid": true,
|
||||
"recall_and_mrr_measured": true,
|
||||
"real_dense_model": true,
|
||||
"index_usage_and_cost_measured": true,
|
||||
"raw_request_response_receipts": true,
|
||||
"all_calls_succeeded": true,
|
||||
"passed": true
|
||||
}
|
||||
}
|
||||
+1526
File diff suppressed because it is too large
Load Diff
+119
@@ -0,0 +1,119 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-10",
|
||||
"run_id": "20260729T203124Z-3_10-476fa870",
|
||||
"created_at": "2026-07-29T20:31:24.286849+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/validation/runs/20260729T203124Z-3_10-476fa870",
|
||||
"artifacts": {
|
||||
"evidence.json": "e47234c3d01b59dbc5acb35899a140f435f9bd491aa442edf1b509d48b5d1e97",
|
||||
"receipts.json": "8cf26248f9fd99206784831992e03d67895e51d92f3a2d4223ed56e22b8a6ca2"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/document_store.json",
|
||||
"sha256": "1fe9c0a7096fc9fcf50c2c5c8ea2008b001243932043ab8882cff436671f8d48",
|
||||
"bytes": 154529
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/evaluation/retrieval_eval.json",
|
||||
"sha256": "7df53f5877ac6a8ec96aa3b2dc575e8cc556e1c9a3eb0ce28800ddd20d0f4a99",
|
||||
"bytes": 4426
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/1-宪法/宪法.md",
|
||||
"sha256": "4cbc5d5d8a4632eb5587ab6377ae5251ce21a149716690864db0b0ab6cd4b386",
|
||||
"bytes": 52540
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/2-宪法相关法/检察官法(2019-04-23).md",
|
||||
"sha256": "16285739d98e9661208bf8b1f334d40d902af794dcef322605bc93ffbfb7c530",
|
||||
"bytes": 20824
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"methods": {
|
||||
"plain_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.6,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.7508333333333334
|
||||
},
|
||||
"contextual_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8555555555555555
|
||||
},
|
||||
"plain_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9222222222222223
|
||||
},
|
||||
"contextual_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.9333333333333333,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9666666666666667
|
||||
},
|
||||
"plain_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.7333333333333333,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8444444444444444
|
||||
},
|
||||
"contextual_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9133333333333333
|
||||
}
|
||||
},
|
||||
"index_time": {
|
||||
"context_generation_ms": 61804.353,
|
||||
"embedding_ms": 365453.495,
|
||||
"usage": {
|
||||
"prompt_tokens": 219888,
|
||||
"completion_tokens": 8846,
|
||||
"total_tokens": 228734
|
||||
},
|
||||
"estimated_cost_usd": 0.033918,
|
||||
"pricing_assumption": {
|
||||
"input_per_million_usd": 0.11,
|
||||
"output_per_million_usd": 1.1
|
||||
}
|
||||
},
|
||||
"errors": 0
|
||||
},
|
||||
"acceptance": {
|
||||
"live_prefix_for_every_chunk": true,
|
||||
"full_source_document_and_target_chunk_in_requests": true,
|
||||
"same_chunks_and_queries": true,
|
||||
"plain_contextual_bm25_dense_hybrid": true,
|
||||
"recall_and_mrr_measured": true,
|
||||
"real_dense_model": true,
|
||||
"index_usage_and_cost_measured": true,
|
||||
"raw_request_response_receipts": true,
|
||||
"all_calls_succeeded": true,
|
||||
"passed": true
|
||||
}
|
||||
}
|
||||
+1762
File diff suppressed because one or more lines are too long
+1526
File diff suppressed because it is too large
Load Diff
+129
@@ -0,0 +1,129 @@
|
||||
{
|
||||
"schema_version": "chapter3-evidence-v1",
|
||||
"experiment": "3-10",
|
||||
"run_id": "20260729T205807Z-3_10-35531780",
|
||||
"created_at": "2026-07-29T20:58:07.546305+00:00",
|
||||
"status": "passed",
|
||||
"run_dir": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/validation/runs/20260729T205807Z-3_10-35531780",
|
||||
"artifacts": {
|
||||
"evidence.json": "e8d2ea403646d0ed0041a5849f38bf16d43481ea8b622a30ae9bdf774b60b11b",
|
||||
"receipts.json": "8cf26248f9fd99206784831992e03d67895e51d92f3a2d4223ed56e22b8a6ca2"
|
||||
},
|
||||
"inputs": [
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/campaign.py",
|
||||
"sha256": "2e61651942f869a2374f24fe40e92aaf332a5bd5abf9c3dbc951fff3a31d9e28",
|
||||
"bytes": 13270
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/compare_retrieval.py",
|
||||
"sha256": "5cb1213c44c94a875e14475c3f371799dd750c00001a9ab36b49034b460f94ef",
|
||||
"bytes": 15372
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/document_store.json",
|
||||
"sha256": "1fe9c0a7096fc9fcf50c2c5c8ea2008b001243932043ab8882cff436671f8d48",
|
||||
"bytes": 154529
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/contextual-retrieval/evaluation/retrieval_eval.json",
|
||||
"sha256": "7df53f5877ac6a8ec96aa3b2dc575e8cc556e1c9a3eb0ce28800ddd20d0f4a99",
|
||||
"bytes": 4426
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/1-宪法/宪法.md",
|
||||
"sha256": "4cbc5d5d8a4632eb5587ab6377ae5251ce21a149716690864db0b0ab6cd4b386",
|
||||
"bytes": 52540
|
||||
},
|
||||
{
|
||||
"path": "/Users/boj/book/ai-agent-book/chapter3/agentic-rag/laws/2-宪法相关法/检察官法(2019-04-23).md",
|
||||
"sha256": "16285739d98e9661208bf8b1f334d40d902af794dcef322605bc93ffbfb7c530",
|
||||
"bytes": 20824
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"methods": {
|
||||
"plain_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.6,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.7508333333333334
|
||||
},
|
||||
"contextual_bm25": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8,
|
||||
"3": 0.8666666666666667,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8555555555555555
|
||||
},
|
||||
"plain_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9222222222222223
|
||||
},
|
||||
"contextual_dense": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.9333333333333333,
|
||||
"3": 1.0,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9666666666666667
|
||||
},
|
||||
"plain_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.7333333333333333,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 0.9333333333333333
|
||||
},
|
||||
"mrr": 0.8444444444444444
|
||||
},
|
||||
"contextual_hybrid": {
|
||||
"n": 15,
|
||||
"recall_at_k": {
|
||||
"1": 0.8666666666666667,
|
||||
"3": 0.9333333333333333,
|
||||
"5": 1.0
|
||||
},
|
||||
"mrr": 0.9133333333333333
|
||||
}
|
||||
},
|
||||
"index_time": {
|
||||
"context_generation_ms": 61804.353,
|
||||
"embedding_ms": 365453.495,
|
||||
"usage": {
|
||||
"prompt_tokens": 219888,
|
||||
"completion_tokens": 8846,
|
||||
"total_tokens": 228734
|
||||
},
|
||||
"estimated_cost_usd": 0.033918,
|
||||
"pricing_assumption": {
|
||||
"input_per_million_usd": 0.11,
|
||||
"output_per_million_usd": 1.1
|
||||
}
|
||||
},
|
||||
"errors": 0
|
||||
},
|
||||
"acceptance": {
|
||||
"live_prefix_for_every_chunk": true,
|
||||
"full_source_document_and_target_chunk_in_requests": true,
|
||||
"same_chunks_and_queries": true,
|
||||
"plain_contextual_bm25_dense_hybrid": true,
|
||||
"recall_and_mrr_measured": true,
|
||||
"real_dense_model": true,
|
||||
"index_usage_and_cost_measured": true,
|
||||
"raw_request_response_receipts": true,
|
||||
"all_calls_succeeded": true,
|
||||
"passed": true
|
||||
}
|
||||
}
|
||||
+1762
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user