liqiang
b119135836
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
319 lines
12 KiB
Python
319 lines
12 KiB
Python
"""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"""
|
||
kb = agent.config.knowledge_base
|
||
active_top_k = kb.offline_top_k if kb.type == KnowledgeBaseType.OFFLINE else kb.local_top_k
|
||
print(f"\n{'='*60}")
|
||
print(f"Agentic RAG System - {mode.capitalize()} Mode")
|
||
print(f"Verbose: {'Enabled' if agent.config.agent.verbose else 'Disabled'} | KB: {kb.type.value} | Top-K: {active_top_k}")
|
||
print(f"{'='*60}")
|
||
print("Type 'quit' or 'exit' to stop")
|
||
print("Type 'clear' to clear conversation history")
|
||
print("Type 'mode' to switch between agentic/non-agentic modes")
|
||
print(f"{'='*60}\n")
|
||
|
||
current_mode = mode
|
||
|
||
while True:
|
||
try:
|
||
user_input = input("\n[USER] > ").strip()
|
||
|
||
if user_input.lower() in ['quit', 'exit']:
|
||
print("\nGoodbye!")
|
||
break
|
||
|
||
if user_input.lower() == 'clear':
|
||
agent.clear_history()
|
||
print("Conversation history cleared.")
|
||
continue
|
||
|
||
if user_input.lower() == 'mode':
|
||
current_mode = "non-agentic" if current_mode == "agentic" else "agentic"
|
||
print(f"Switched to {current_mode} mode")
|
||
continue
|
||
|
||
if not user_input:
|
||
continue
|
||
|
||
# Process query
|
||
print(f"\n[ASSISTANT ({current_mode})] > ", end="", flush=True)
|
||
|
||
if current_mode == "agentic":
|
||
response = agent.query(user_input, stream=True)
|
||
else:
|
||
response = agent.query_non_agentic(user_input, stream=True)
|
||
|
||
# Handle streaming response
|
||
if hasattr(response, '__iter__'):
|
||
for chunk in response:
|
||
print(chunk, end="", flush=True)
|
||
print() # New line after response
|
||
else:
|
||
print(response)
|
||
|
||
except KeyboardInterrupt:
|
||
print("\n\nInterrupted. Type 'quit' to exit.")
|
||
except Exception as e:
|
||
logger.error(f"Error: {e}")
|
||
print(f"\nError processing query: {e}")
|
||
|
||
|
||
def run_batch_mode(agent: AgenticRAG, queries_file: str, output_file: str, mode: str = "agentic"):
|
||
"""Run batch queries from file"""
|
||
try:
|
||
with open(queries_file, 'r', encoding='utf-8') as f:
|
||
queries = [line.strip() for line in f if line.strip()]
|
||
except Exception as e:
|
||
logger.error(f"Error reading queries file: {e}")
|
||
return
|
||
|
||
results = []
|
||
|
||
for i, query in enumerate(queries, 1):
|
||
print(f"\n[{i}/{len(queries)}] Processing: {query[:100]}...")
|
||
|
||
try:
|
||
if mode == "agentic":
|
||
response = agent.query(query, stream=False)
|
||
else:
|
||
response = agent.query_non_agentic(query, stream=False)
|
||
|
||
results.append({
|
||
"query": query,
|
||
"response": response,
|
||
"mode": mode
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error processing query: {e}")
|
||
results.append({
|
||
"query": query,
|
||
"response": f"Error: {str(e)}",
|
||
"mode": mode
|
||
})
|
||
|
||
# Save results
|
||
try:
|
||
with open(output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||
print(f"\nResults saved to {output_file}")
|
||
except Exception as e:
|
||
logger.error(f"Error saving results: {e}")
|
||
|
||
|
||
def run_comparison_mode(agent: AgenticRAG, query: str):
|
||
"""Run both modes and compare results"""
|
||
print(f"\n{'='*60}")
|
||
print("Comparison Mode - Running both Agentic and Non-Agentic")
|
||
print(f"{'='*60}")
|
||
print(f"Query: {query}")
|
||
print(f"{'='*60}")
|
||
|
||
# Run non-agentic mode
|
||
print("\n[NON-AGENTIC MODE]")
|
||
print("-" * 40)
|
||
non_agentic_response = agent.query_non_agentic(query, stream=False)
|
||
print(non_agentic_response)
|
||
|
||
# Clear history for fair comparison
|
||
agent.clear_history()
|
||
|
||
# Run agentic mode
|
||
print("\n[AGENTIC MODE]")
|
||
print("-" * 40)
|
||
agentic_response = agent.query(query, stream=False)
|
||
print(agentic_response)
|
||
|
||
print(f"\n{'='*60}")
|
||
|
||
|
||
def main():
|
||
"""Main function"""
|
||
parser = argparse.ArgumentParser(
|
||
description="智能体化 RAG 系统：对比『智能体化（多轮迭代检索）』与『非智能体化（单次检索）』两种范式。",
|
||
epilog=(
|
||
"示例:\n"
|
||
" python main.py --kb-type offline --query \"醉酒过失致人重伤且有盗窃前科如何量刑\"\n"
|
||
" python main.py --query \"故意杀人罪判几年\" --mode compare --kb-type offline\n"
|
||
" python compare_offline.py # 纯离线检索对比，无需 API 与外部服务\n"
|
||
),
|
||
formatter_class=argparse.RawTextHelpFormatter,
|
||
)
|
||
|
||
# 模式选择
|
||
parser.add_argument("--mode", choices=["agentic", "non-agentic", "compare"],
|
||
default="agentic",
|
||
help="查询模式：agentic=智能体化多轮检索 / non-agentic=单次检索 / compare=同题对比（默认：agentic）")
|
||
|
||
# 查询选项
|
||
parser.add_argument("--query", type=str, help="单条查询问题；不指定则进入交互模式")
|
||
parser.add_argument("--batch", type=str, help="批量查询文件路径（每行一个问题）")
|
||
parser.add_argument("--output", type=str, default="results.json",
|
||
help="批量结果的输出文件路径（默认：results.json）")
|
||
|
||
# 配置选项
|
||
parser.add_argument("--provider", type=str, help="LLM 提供商（如 dashscope/qwen/bailian / kimi / doubao / openai）")
|
||
parser.add_argument("--model", type=str, help="LLM 模型名称（不指定则用提供商默认模型）")
|
||
parser.add_argument("--kb-type", choices=["offline", "local", "dify"],
|
||
help="知识库后端：offline=内置离线 BM25（无需服务/无需 API）/ local=检索流水线服务 / dify=Dify API")
|
||
parser.add_argument("--corpus", type=str,
|
||
help="离线后端的法律语料目录（仅 --kb-type offline 生效，默认：laws）")
|
||
parser.add_argument("--top-k", type=int, dest="top_k",
|
||
help="检索深度：每次检索返回的分块数量（默认：offline=5，local=3）")
|
||
parser.add_argument("--verbose", action="store_true", help="输出详细的 Agent 推理轨迹（默认开启）")
|
||
parser.add_argument("--no-verbose", action="store_true", help="关闭详细日志输出")
|
||
|
||
# 索引选项
|
||
parser.add_argument("--index", type=str, help="待索引的文件或目录路径")
|
||
parser.add_argument("--chunk-size", type=int, default=2048, help="索引时的分块大小（字符数，默认：2048）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Setup environment
|
||
if not setup_environment():
|
||
logger.warning("Environment setup incomplete, continuing anyway...")
|
||
|
||
# Load or create config
|
||
config = Config.from_env()
|
||
|
||
# Set verbose mode by default (can be disabled with --no-verbose)
|
||
config.agent.verbose = True # Default to verbose mode
|
||
|
||
# Override config with command line args
|
||
if args.provider:
|
||
config.llm.provider = args.provider
|
||
if args.model:
|
||
config.llm.model = args.model
|
||
if args.kb_type:
|
||
config.knowledge_base.type = KnowledgeBaseType(args.kb_type)
|
||
if args.corpus:
|
||
config.knowledge_base.offline_corpus_path = args.corpus
|
||
if args.top_k:
|
||
# 同时设置离线与本地后端的检索深度，保持行为一致
|
||
config.knowledge_base.offline_top_k = args.top_k
|
||
config.knowledge_base.local_top_k = args.top_k
|
||
|
||
# Handle verbose mode (default is True, can be disabled with --no-verbose)
|
||
if args.no_verbose:
|
||
config.agent.verbose = False
|
||
elif args.verbose:
|
||
config.agent.verbose = True # Explicitly set if --verbose is used
|
||
|
||
# Handle indexing if requested
|
||
if args.index:
|
||
print(f"\n{'='*60}")
|
||
print("Indexing Documents")
|
||
print(f"{'='*60}")
|
||
|
||
config.chunking.chunk_size = args.chunk_size
|
||
indexer = DocumentIndexer(config.knowledge_base, config.chunking)
|
||
|
||
from pathlib import Path
|
||
path = Path(args.index)
|
||
|
||
if path.is_file():
|
||
result = indexer.index_file(str(path))
|
||
elif path.is_dir():
|
||
result = indexer.index_directory(str(path))
|
||
else:
|
||
print(f"Path not found: {path}")
|
||
return
|
||
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||
print(f"{'='*60}\n")
|
||
|
||
# Create agent
|
||
agent = AgenticRAG(config)
|
||
|
||
# Handle different execution modes
|
||
if args.query and args.mode == "compare":
|
||
# Comparison mode with single query
|
||
run_comparison_mode(agent, args.query)
|
||
|
||
elif args.query:
|
||
# Single query mode
|
||
kb = config.knowledge_base
|
||
active_top_k = kb.offline_top_k if kb.type == KnowledgeBaseType.OFFLINE else kb.local_top_k
|
||
print(f"\n[Query] {args.query}")
|
||
print(f"[Mode] {args.mode}")
|
||
print(f"[KB] {kb.type.value}")
|
||
print(f"[Verbose] {'Enabled' if config.agent.verbose else 'Disabled'}")
|
||
print(f"[Top-K] {active_top_k}")
|
||
print("-" * 40)
|
||
|
||
if args.mode == "agentic":
|
||
response = agent.query(args.query, stream=False)
|
||
else:
|
||
response = agent.query_non_agentic(args.query, stream=False)
|
||
|
||
print(response)
|
||
|
||
elif args.batch:
|
||
# Batch mode
|
||
run_batch_mode(agent, args.batch, args.output, args.mode)
|
||
|
||
else:
|
||
# Interactive mode (default)
|
||
run_interactive_mode(agent, args.mode)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|