RAG实战（三）路由与查询构建 📍 置顶

from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
import asyncio
from concurrent.futures import ThreadPoolExecutor
import time

# 初始化嵌入模型
embeddings = HuggingFaceEmbeddings(
    model_name="./Models/maidalun/bce-embedding-base_v1",
    model_kwargs={"device": "cuda"},
    encode_kwargs={"normalize_embeddings": True}
)

# 创建 web_content 集合
vectorstore = Chroma(
    collection_name="web_content",  # 指定集合名称
    persist_directory="./chroma_db",
    embedding_function=embeddings
)

print("✅ 创建 web_content 集合成功")

def load_single_url(url: str):
    """加载单个URL"""
    try:
        loader = WebBaseLoader(url)
        docs = loader.load()
        print(f"✅ 成功加载: {url} -> {len(docs)} 个文档")
        return docs
    except Exception as e:
        print(f"❌ 加载失败: {url} -> 错误: {e}")
        return []

async def load_urls_parallel(urls: list, max_workers: int = 5):
    """并行加载多个URL"""
    print(f"🚀 开始并行加载 {len(urls)} 个URL...")
    
    loop = asyncio.get_event_loop()
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 创建并行任务
        tasks = [
            loop.run_in_executor(executor, load_single_url, url)
            for url in urls
        ]
        
        # 并行执行
        all_docs = await asyncio.gather(*tasks)
    
    # 合并所有文档
    flat_docs = []
    for docs in all_docs:
        flat_docs.extend(docs)
    
    print(f"📊 总共加载了 {len(flat_docs)} 个文档")
    return flat_docs

def add_docs_to_vectorstore(docs: list):
    """将文档添加到向量数据库"""
    if not docs:
        print("⚠️ 没有文档可添加")
        return
    
    print("💾 正在将文档添加到向量数据库...")
    
    # 添加文档到集合
    vectorstore.add_documents(docs)
    
    # 持久化保存
    vectorstore.persist()
    
    print(f"✅ 成功添加 {len(docs)} 个文档到 web_content 集合")

# 使用示例
async def main():
    # URL列表
    url_list = [
        "https://blog.csdn.net/qq_40081208/article/details/111053208",
        "https://blog.csdn.net/Yyuan12345678/article/details/142108850",
        "https://blog.csdn.net/WhiffeYF/article/details/111031270",
        "https://www.hanspub.org/journal/PaperInformation?paperID=84081",
        "https://blog.csdn.net/WhiffeYF/article/details/110829105"
    ]
    
    if not url_list:
        print("⚠️ 请先在 url_list 中添加URL链接")
        return
    
    # 1. 并行加载所有URL
    all_docs = await load_urls_parallel(url_list)
    
    # 2. 添加到向量数据库
    add_docs_to_vectorstore(all_docs)
    
    # 3. 验证添加结果
    print("\n🔍 验证数据库内容:")
    collection_info = vectorstore._client.get_collection("web_content")
    print(f"集合中的文档数量: {collection_info.count()}")
    
    # 4. 测试检索
    test_query = "抓取检测"
    results = vectorstore.similarity_search(test_query, k=2)
    print(f"\n🔎 测试检索 '{test_query}':")
    for i, doc in enumerate(results, 1):
        print(f"   {i}. 来源: {doc.metadata.get('source', 'N/A')}")
        print(f"      内容预览: {doc.page_content[:100]}...")

await main()

✅ 创建 web_content 集合成功
🚀 开始并行加载 5 个URL...
✅ 成功加载: https://www.hanspub.org/journal/PaperInformation?paperID=84081 -> 1 个文档
✅ 成功加载: https://blog.csdn.net/WhiffeYF/article/details/110829105 -> 1 个文档
✅ 成功加载: https://blog.csdn.net/Yyuan12345678/article/details/142108850 -> 1 个文档
✅ 成功加载: https://blog.csdn.net/WhiffeYF/article/details/111031270 -> 1 个文档
✅ 成功加载: https://blog.csdn.net/qq_40081208/article/details/111053208 -> 1 个文档
📊 总共加载了 5 个文档
💾 正在将文档添加到向量数据库...
✅ 成功添加 5 个文档到 web_content 集合

🔍 验证数据库内容:
集合中的文档数量: 5

🔎 测试检索 '抓取检测':
   1. 来源: https://blog.csdn.net/qq_40081208/article/details/111053208
      内容预览: 抓取检测之Dex-Net 2.0_dexnet-CSDN博客...
   2. 来源: https://blog.csdn.net/WhiffeYF/article/details/111031270
      内容预览: 械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状...

from vllm import SamplingParams

# 定义路由提示词（ChatML格式）
route_prompt_template = """<|im_start|>system
		你是一个路由助手，负责将用户查询发送到正确的数据源。
		
		可用的数据源:
		- langchain: 学术论文和技术文档，包含抓取检测、滑动检测相关的研究论文、技术文档
		- web_content: 网页内容，包含CSDN博客、技术教程、实践指南等网页文章
		
		请分析用户查询，返回最合适的数据源名称(只返回名称，不要其他内容)。<|im_end|>
		<|im_start|>user
		{question}<|im_end|>
		<|im_start|>assistant
		数据源："""

def route_chain(question: str, llm) -> str:
    """基础路由链"""
    prompt = route_prompt_template.format(question=question)
    
    # 使用 vLLM 的正确调用方式
    sampling_params = SamplingParams(
        temperature=0.1,
        top_p=0.9,
        max_tokens=50,
        stop=["<|im_end|>", "<|endoftext|>"]
    )
    
    # 生成响应
    outputs = llm.generate([prompt], sampling_params)
    
    # 提取结果
    if outputs and len(outputs) > 0:
        output = outputs[0]
        if hasattr(output, 'outputs') and output.outputs:
            return output.outputs[0].text.strip()
    
    return "web_content"  # 默认返回

# 测试路由
question1 = "如何在数据库中查询？"
route1 = route_chain(question1, llm)
print(f"Query: {question1}")
print(f"Route: {route1}\n")

question2 = "今天的天气怎么样？"
route2 = route_chain(question2, llm)
print(f"Query: {question2}")
print(f"Route: {route2}\n")

from langchain.vectorstores import Chroma
from typing import Dict, Any, List

class LogicalRouter:
    """完整的逻辑路由器"""
    
    def __init__(self, llm, embeddings, persist_directory="./chroma_db"):
        self.llm = llm
        self.embeddings = embeddings
        self.persist_directory = persist_directory
        self.routes = {}
        self.route_descriptions = {}
        
        # 初始化集合
        self._init_collections()
        print("✅ 逻辑路由器初始化完成")
    
    def _init_collections(self):
        """初始化向量数据库集合"""
        # 1. langchain 集合（论文）
        self.langchain_collection = Chroma(
            collection_name="langchain",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        
        # 2. web_content 集合（网页内容）
        self.web_content_collection = Chroma(
            collection_name="web_content",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        
        # 自动添加路由
        self.add_route("langchain", self.langchain_collection.as_retriever(), 
                      "学术论文和技术文档，包含抓取检测、滑动检测相关的研究论文、技术文档")
        self.add_route("web_content", self.web_content_collection.as_retriever(), 
                      "网页内容，包含CSDN博客、技术教程、实践指南等网页文章")
    
    def add_route(self, name: str, retriever, description: str = None):
        """添加路由"""
        self.routes[name] = retriever
        if description:
            self.route_descriptions[name] = description
    
    def _call_llm(self, prompt: str, max_tokens: int = 100) -> str:
        """调用LLM进行路由决策"""
        sampling_params = SamplingParams(
            temperature=0.1,
            top_p=0.9,
            max_tokens=max_tokens,
            stop=["<|im_end|>", "<|endoftext|>"]
        )
        
        outputs = self.llm.generate([prompt], sampling_params)
        
        if outputs and outputs[0].outputs:
            return outputs[0].outputs[0].text.strip()
        return ""
    
    def route(self, question: str) -> tuple:
        """执行路由决策"""
        # 构建路由提示词
        route_descriptions = self._get_route_descriptions()
        
        prompt = f"""<|im_start|>system
            你是一个路由助手，负责将用户查询发送到正确的数据源。
            
            可用的数据源：
            {route_descriptions}
            
            请分析用户查询，返回最合适的数据源名称（只返回名称，不要其他内容）。<|im_end|>
            <|im_start|>user
            问题：{question}<|im_end|>
            <|im_start|>assistant
            数据源："""
        
        # 调用LLM进行路由决策
        route_name = self._call_llm(prompt, max_tokens=50)
        route_name = route_name.strip().lower()
        
        print(f"🤖 LLM路由决策: '{route_name}'")
        
        # 检查路由是否存在，支持模糊匹配
        if route_name not in self.routes:
            for name in self.routes.keys():
                if name in route_name or route_name in name:
                    route_name = name
                    print(f"🔄 模糊匹配到: {route_name}")
                    break
            else:
                # 默认回退
                print(f"⚠️ 未找到路由: {route_name}，使用默认路由: web_content")
                route_name = "web_content"
        
        return route_name, self.routes[route_name]
    
    def _get_route_descriptions(self) -> str:
        """获取路由描述"""
        descriptions = []
        for name, retriever in self.routes.items():
            desc = self.route_descriptions.get(name, f'{name}数据源')
            descriptions.append(f"- {name}: {desc}")
        return "\n".join(descriptions)
    
    def query(self, question: str, k: int = 4) -> Dict[str, Any]:
        """执行完整查询"""
        print(f"🎯 查询: {question}")
        
        # 路由到正确的数据源
        route_name, retriever = self.route(question)
        print(f"📍 路由到: {route_name}")
        
        # 执行检索
        docs = retriever.get_relevant_documents(question, k=k)
        print(f"📚 检索到 {len(docs)} 个文档")
        
        return {
            "route": route_name,
            "documents": docs,
            "question": question,
            "collection_size": self._get_collection_size(route_name)
        }
    
    def _get_collection_size(self, collection_name: str) -> int:
        """获取集合大小"""
        try:
            if collection_name == "langchain":
                collection = self.langchain_collection._client.get_collection("langchain")
            else:
                collection = self.web_content_collection._client.get_collection("web_content")
            return collection.count()
        except:
            return 0
    
    def get_collection_info(self) -> Dict[str, Any]:
        """获取集合信息"""
        info = {}
        for name in ["langchain", "web_content"]:
            try:
                if name == "langchain":
                    collection = self.langchain_collection._client.get_collection("langchain")
                else:
                    collection = self.web_content_collection._client.get_collection("web_content")
                info[name] = {
                    "document_count": collection.count(),
                    "description": self.route_descriptions.get(name, 'N/A')
                }
            except Exception as e:
                info[name] = {"error": str(e)}
        
        return info

# 初始化路由器
router = LogicalRouter(llm, embeddings, persist_directory="./chroma_db")

# 查看集合信息
print("📊 集合信息:")
info = router.get_collection_info()
for name, data in info.items():
    if "document_count" in data:
        print(f"   {name}: {data['document_count']} 个文档 - {data['description']}")

# 测试路由查询
print("\n" + "="*50)
result = router.query("网页文档中关于抓取检测的介绍")
print(f"路由结果: {result['route']}")
print(f"文档数量: {len(result['documents'])}")

# 显示文档预览
for i, doc in enumerate(result['documents'][:2], 1):
    source = doc.metadata.get('source', '未知来源')
    print(f"  {i}. 来源: {source}")

class LogicalRouterWithFallback:
    """带回退机制的逻辑路由器"""
    
    def __init__(self, llm, embeddings, fallback_route: str = "web_content", persist_directory="./chroma_db"):
        self.llm = llm
        self.embeddings = embeddings
        self.persist_directory = persist_directory
        self.fallback_route = fallback_route
        self.routes = {}
        self.route_descriptions = {}
        
        # 初始化集合
        self._init_collections()
        print("✅ 带回退机制的逻辑路由器初始化完成")
    
    def _init_collections(self):
        """初始化集合（与之前相同）"""
        self.langchain_collection = Chroma(
            collection_name="langchain",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        
        self.web_content_collection = Chroma(
            collection_name="web_content",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        
        # 添加路由
        self.add_route("langchain", self.langchain_collection.as_retriever(), 
                      "学术论文和技术文档")
        self.add_route("web_content", self.web_content_collection.as_retriever(), 
                      "网页内容")
    
    def add_route(self, name: str, retriever, description: str = None):
        """添加路由"""
        self.routes[name] = retriever
        if description:
            self.route_descriptions[name] = description
    
    def route(self, question: str) -> tuple:
        """执行路由，带回退机制"""
        try:
            # 构建路由提示词
            route_descriptions = self._get_route_descriptions()
            
            prompt = f"""<|im_start|>system
                你是一个路由助手，负责将用户查询发送到正确的数据源。
                
                可用的数据源：
                {route_descriptions}
                
                请分析用户查询，返回最合适的数据源名称（只返回名称，不要其他内容）。<|im_end|>
                <|im_start|>user
                问题：{question}<|im_end|>
                <|im_start|>assistant
                数据源："""
            
            # 调用LLM
            route_name = self._call_llm(prompt, max_tokens=50)
            route_name = route_name.strip().lower()
            
            print(f"🤖 LLM路由决策: '{route_name}'")
            
            # 检查路由是否存在
            if route_name not in self.routes:
                for name in self.routes.keys():
                    if name in route_name or route_name in name:
                        route_name = name
                        print(f"🔄 模糊匹配到: {route_name}")
                        break
                else:
                    raise ValueError(f"未找到路由: {route_name}")
            
            return route_name, self.routes[route_name]
            
        except Exception as e:
            print(f"⚠️ 路由失败: {e}, 使用回退路由: {self.fallback_route}")
            return self.fallback_route, self.routes[self.fallback_route]
    
    def _call_llm(self, prompt: str, max_tokens: int = 100) -> str:
        """调用LLM"""
        sampling_params = SamplingParams(
            temperature=0.1,
            top_p=0.9,
            max_tokens=max_tokens,
            stop=["<|im_end|>", "<|endoftext|>"]
        )
        
        outputs = self.llm.generate([prompt], sampling_params)
        
        if outputs and outputs[0].outputs:
            return outputs[0].outputs[0].text.strip()
        return ""
    
    def query_multiple(self, question: str, max_routes: int = 2, k: int = 3) -> Dict[str, Any]:
        """查询多个数据源（带回退机制）"""
        print(f"🎯 多数据源查询: {question}")
        
        # 获取主路由
        primary_route, primary_retriever = self.route(question)
        print(f"📍 主路由: {primary_route}")
        
        # 获取主路由文档
        primary_docs = primary_retriever.get_relevant_documents(question, k=k)
        results = {
            primary_route: primary_docs
        }
        print(f"📚 {primary_route}: {len(primary_docs)} 个文档")
        
        # 如果需要，添加回退路由
        if len(results) < max_routes and primary_route != self.fallback_route:
            print(f"🔄 添加回退路由: {self.fallback_route}")
            fallback_docs = self.routes[self.fallback_route].get_relevant_documents(question, k=k)
            results[self.fallback_route] = fallback_docs
            print(f"📚 {self.fallback_route}: {len(fallback_docs)} 个文档")
        
        return results
    
    def _get_route_descriptions(self) -> str:
        """获取路由描述"""
        descriptions = []
        for name, retriever in self.routes.items():
            desc = self.route_descriptions.get(name, f'{name}数据源')
            descriptions.append(f"- {name}: {desc}")
        return "\n".join(descriptions)

# 创建带回退机制的路由器
router = LogicalRouterWithFallback(
    llm=llm,
    embeddings=embeddings,
    fallback_route="web_content",
    persist_directory="./chroma_db"
)

# 多数据源查询示例
query = "论文中抓取检测的定义"
results = router.query_multiple(query, max_routes=2, k=3)

for route_name, docs in results.items():
    print(f"\n{route_name}: {len(docs)} 个文档")
    for i, doc in enumerate(docs[:2], 1):
        source = doc.metadata.get('source', '未知来源')
        print(f"  {i}. 来源: {source}")
        clean_content = doc.page_content[:80].replace('\n', ' ').replace('\r', '').strip()
        print(f"     内容: {clean_content}...")

🎯 多数据源查询: 论文中抓取检测的定义
Processed prompts: 100%|██████████| 1/1 [00:00<00:00, 12.62it/s]
🤖 LLM路由决策: 'langchain'
📍 主路由: langchain
📚 langchain: 4 个文档
🔄 添加回退路由: web_content
📚 web_content: 4 个文档

langchain: 4 个文档
  1. 来源: ./Dataset/PDF/基于视触觉融合的机器人物体识别和抓取稳定性检测的研究与应用_上官明雨.pdf
     内容: 1.4 本论文组织结构  本文总体框架图如图 1.4 所示。本文将研究内容分为六个章节进行阐述：  第一章，绪论。本章主要介绍了课题的研究背景与研究意义，分析了...
  2. 来源: ./Dataset/PDF/基于视触感知协同的机器人抓取技术研究_祝会龙.pdf
     内容: 西南科技大学硕士学位论文  32    图 4-2 抓取过程   Fig.4-2 Grabbing Process   （1）抓手打开阶段：抓手在初始位置 将抓...

web_content: 4 个文档
  1. 来源: https://blog.csdn.net/WhiffeYF/article/details/111031270
     内容: 械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状 刘亚欣_基于深度图像的机械臂抓取位姿估计与轨迹优化研究-CSDN博...
  2. 来源: https://blog.csdn.net/qq_40081208/article/details/111053208
     内容: 抓取检测之Dex-Net 2.0_dexnet-CSDN博客...

class OptimizedLogicalRouter(LogicalRouterWithFallback):
    """优化版逻辑路由器"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.route_cache = {}  # 路由缓存
        self.cache_size = 1000  # 缓存大小
    
    def _get_cache_key(self, question: str) -> str:
        """生成缓存键"""
        import hashlib
        return hashlib.md5(question.encode()).hexdigest()
    
    def route(self, question: str) -> tuple:
        """带缓存的路由"""
        cache_key = self._get_cache_key(question)
        
        # 检查缓存
        if cache_key in self.route_cache:
            print("💾 使用缓存路由")
            return self.route_cache[cache_key]
        
        # 执行路由
        result = super().route(question)
        
        # 更新缓存
        if len(self.route_cache) >= self.cache_size:
            # 简单LRU策略：移除第一个
            first_key = next(iter(self.route_cache))
            del self.route_cache[first_key]
        
        self.route_cache[cache_key] = result
        return result
    
    def batch_route(self, questions: List[str]) -> List[tuple]:
        """批量路由"""
        results = []
        for question in questions:
            result = self.route(question)
            results.append((question, result))
        return results

import numpy as np
from typing import Dict, List, Tuple
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

class SemanticRouter:
    """语义路由器（适配本地配置）"""
    
    def __init__(self, embeddings, persist_directory="./chroma_db"):
        self.embeddings = embeddings
        self.persist_directory = persist_directory
        self.routes = {}
        self.route_embeddings = {}
        
        # 初始化向量集合
        self._init_collections()
    
    def _init_collections(self):
        """初始化向量数据库集合"""
        # 学术论文集合
        self.langchain_collection = Chroma(
            collection_name="langchain",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
        
        # 网页内容集合
        self.web_content_collection = Chroma(
            collection_name="web_content",
            persist_directory=self.persist_directory,
            embedding_function=self.embeddings
        )
    
    def add_route(self, name: str, description: str, retriever):
        """添加路由规则"""
        self.routes[name] = {
            'description': description,
            'retriever': retriever
        }
        
        # 预计算路由描述的嵌入向量
        self.route_embeddings[name] = self.embeddings.embed_query(description)
        print(f"✅ 添加路由: {name} - {description}")
    
    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
        """计算余弦相似度"""
        vec1 = np.array(vec1)
        vec2 = np.array(vec2)
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        return dot_product / (norm1 * norm2) if norm1 and norm2 else 0.0
    
    def route(self, question: str, threshold: float = 0.3) -> Tuple[str, float]:
        """执行语义路由"""
        query_embedding = self.embeddings.embed_query(question)
        
        similarities = {}
        for name, route_embedding in self.route_embeddings.items():
            similarity = self._cosine_similarity(query_embedding, route_embedding)
            similarities[name] = similarity
        
        best_route = max(similarities, key=similarities.get)
        best_score = similarities[best_route]
        
        print(f"🔍 路由分析:")
        for route, score in sorted(similarities.items(), key=lambda x: x[1], reverse=True):
            print(f"   {route}: {score:.3f}")
        
        return best_route, best_score
    
    def query(self, question: str, k: int = 4):
        """完整查询流程"""
        print(f"🎯 查询: {question}")
        
        # 语义路由
        route_name, score = self.route(question)
        print(f"📍 路由到: {route_name} (相似度: {score:.3f})")
        
        # 检索
        retriever = self.routes[route_name]['retriever']
        docs = retriever.get_relevant_documents(question, k=k)
        
        return {
            "route": route_name,
            "score": score,
            "documents": docs
        }

def query_multiple(self, question: str, k: int = 3):
    """多数据源并行查询"""
    scores = self.route_with_scores(question)
    
    results = {}
    for route_name, score in scores.items():
        if score > 0.2:  # 相似度阈值
            retriever = self.routes[route_name]['retriever']
            docs = retriever.get_relevant_documents(question, k=k)
            results[route_name] = {
                'documents': docs,
                'score': score
            }
    
    return results

# 使用示例
multi_result = router.query_multiple("抓取检测的最新进展")

用户查询: "如何学习抓取检测算法？"
    ↓
自适应路由分析
    ↓
复杂度判断: simple (简单查询)
    ↓
选择语义路由策略
    ↓
语义路由流程开始
    ↓
计算查询嵌入
    ↓
计算与各路由描述的相似度
    ↓
选择最相似的路由: langchain (分数: 0.45)
    ↓
执行检索: 从langchain集合检索文档
    ↓
返回结果: 4个相关文档

from typing import Dict, List, Any
import numpy as np

class HybridRouter:
    """混合路由器：结合逻辑和语义路由（使用现有实现）"""
    
    def __init__(self, embeddings, llm, persist_directory="./chroma_db"):
        self.embeddings = embeddings
        self.llm = llm
        self.persist_directory = persist_directory
        self.routes = {}
        
        # 使用之前实现的语义路由器和逻辑路由器
        self.semantic_router = SemanticRouter(embeddings, persist_directory)
        self.logical_router = LogicalRouter(llm, embeddings, persist_directory)
    
    def add_route(self, name: str, description: str, retriever):
        """添加路由到两个路由器"""
        # 添加到语义路由器
        self.semantic_router.add_route(name, description, retriever)
        # 添加到逻辑路由器
        self.logical_router.add_route(name, retriever)
        # 添加到本地路由表
        self.routes[name] = retriever
        print(f"✅ 添加混合路由: {name} - {description}")
    
    def route(self, question: str, use_semantic: bool = True, semantic_threshold: float = 0.3):
        """执行混合路由
        
        Args:
            question: 用户查询
            use_semantic: 是否优先使用语义路由
            semantic_threshold: 语义路由阈值（本地模型阈值较低）
        """
        print(f"🎯 混合路由查询: {question}")
        
        if use_semantic:
            # 尝试语义路由
            try:
                route_name, score = self.semantic_router.route(question, threshold=semantic_threshold)
                
                if score >= semantic_threshold:
                    print(f"✅ 使用语义路由: {route_name} (分数: {score:.3f})")
                    return route_name, self.routes[route_name]
                else:
                    print(f"⚠️ 语义路由分数过低 ({score:.3f} < {semantic_threshold})，切换到逻辑路由")
            except Exception as e:
                print(f"⚠️ 语义路由失败: {e}，切换到逻辑路由")
        
        # 使用逻辑路由作为后备
        try:
            route_name, retriever = self.logical_router.route(question)
            print(f"✅ 使用逻辑路由: {route_name}")
            return route_name, retriever
        except Exception as e:
            print(f"❌ 逻辑路由失败: {e}，使用默认路由")
            # 回退到web_content
            return "web_content", self.routes["web_content"]
    
    def query(self, question: str, use_semantic: bool = True, k: int = 4):
        """执行查询"""
        route_name, retriever = self.route(question, use_semantic)
        docs = retriever.get_relevant_documents(question, k=k)
        
        print(f"📚 检索到 {len(docs)} 个文档")
        
        return {
            "route": route_name,
            "documents": docs,
            "question": question
        }
    
    def query_adaptive(self, question: str, k: int = 4):
        """自适应查询：根据查询复杂度选择最佳路由"""
        print(f"🎯 自适应查询: {question}")
        
        # 分析查询复杂度
        complexity = self._analyze_complexity(question)
        print(f"📊 查询复杂度: {complexity}")
        
        if complexity == "simple":
            # 简单查询：使用语义路由
            return self.query(question, use_semantic=True, k=k)
        else:
            # 复杂查询：使用逻辑路由
            return self.query(question, use_semantic=False, k=k)
    
    def _analyze_complexity(self, question: str) -> str:
        """分析查询复杂度"""
        # 简单规则：根据查询长度和关键词判断
        if len(question) < 20 and any(keyword in question.lower() for keyword in ["是什么", "怎么用", "如何", "教程"]):
            return "simple"
        else:
            return "complex"

from langchain.chains.query_constructor.base import AttributeInfo
from vllm import SamplingParams
import re

class CustomSelfQueryRetriever:
    """自定义自查询检索器（适配vLLM配置）"""
    
    def __init__(self, llm, vectorstore, metadata_field_info, document_content_description):
        self.llm = llm
        self.vectorstore = vectorstore
        self.metadata_field_info = metadata_field_info
        self.document_content_description = document_content_description
    
    def _call_llm(self, prompt: str, max_tokens: int = 200) -> str:
        """调用vLLM模型进行查询解析"""
        sampling_params = SamplingParams(
            temperature=0.1,
            top_p=0.9,
            max_tokens=max_tokens,
            stop=["<|im_end|>", "<|endoftext|>"]
        )
        
        outputs = self.llm.generate([prompt], sampling_params)
        if outputs and outputs[0].outputs:
            return outputs[0].outputs[0].text.strip()
        return ""
    
    def _parse_filter_query(self, query: str) -> dict:
        """解析自然语言查询，生成过滤条件"""
        # 构建元数据字段描述
        metadata_fields = "\n".join([
            f"- {field.name}: {field.description} (类型: {field.type})" 
            for field in self.metadata_field_info
        ])
        
        prompt = f"""<|im_start|>system
					你是一个查询解析器，负责将自然语言查询转换为结构化过滤条件。
					
					可用的元数据字段：
					{metadata_fields}
					
					文档内容描述：{self.document_content_description}
					
					请将用户查询解析为过滤条件，格式为：
					filter_type:field_name:value
					
					支持的过滤操作：
					- eq: 等于
					- gt: 大于  
					- lt: 小于
					- contains: 包含（用于内容语义搜索）
					
					示例：
					查询: "找出2023年关于Python的文章"
					解析: eq:date:2023;contains:content:Python
					
					查询: "阅读量超过1000的技术文章"  
					解析: gt:views:1000;contains:content:技术
					
					只返回解析后的过滤条件，不要其他内容。<|im_end|>
					<|im_start|>user
					查询：{query}<|im_end|>
					<|im_start|>assistant
					解析："""
        
        response = self._call_llm(prompt)
        print(f"🤖 LLM解析结果: {response}")
        return self._parse_filter_response(response)
    
    def _parse_filter_response(self, response: str) -> dict:
        """解析LLM返回的过滤条件"""
        filters = {}
        
        # 解析格式: eq:date:2023;contains:content:Python
        filter_parts = response.split(';')
        
        for part in filter_parts:
            part = part.strip()
            if part.count(':') >= 2:
                try:
                    # 拆分成 [操作符, 字段名, 值]
                    op, field, value = part.split(':', 2)
                    if field not in filters:
                        filters[field] = {}
                    filters[field][op] = value
                except ValueError:
                    continue
        
        return filters
    
    def get_relevant_documents(self, query: str, k: int = 4):
        """获取相关文档（核心方法）"""
        print(f"🎯 自查询: {query}")
        
        try:
            # 1. 解析查询生成过滤条件
            filters = self._parse_filter_query(query)
            print(f"🔍 解析的过滤条件: {filters}")
            
            # 2. 构建过滤条件
            where_clauses = []
            search_query = query  # 默认搜索词
            
            for field, conditions in filters.items():
                for op, value in conditions.items():
                    if op == 'eq':
                        where_clauses.append({field: {"$eq": value}})
                    elif op == 'gt':
                        where_clauses.append({field: {"$gt": int(value)}})
                    elif op == 'lt':
                        where_clauses.append({field: {"$lt": int(value)}})
                    elif op == 'contains' and field == 'content':
                        # 使用解析出的内容作为搜索词
                        search_query = value
            
            # 3. 执行检索
            if where_clauses:
                # 组合过滤条件
                if len(where_clauses) == 1:
                    combined_filter = where_clauses[0]
                else:
                    combined_filter = {"$and": where_clauses}
                
                docs = self.vectorstore.similarity_search(
                    search_query, 
                    k=k, 
                    filter=combined_filter
                )
            else:
                # 无过滤条件，直接搜索
                docs = self.vectorstore.similarity_search(search_query, k=k)
            
            return docs
            
        except Exception as e:
            print(f"❌ 自查询失败: {e}，使用普通检索")
            # 回退到普通检索
            return self.vectorstore.similarity_search(query, k=k)

# 定义元数据字段信息
metadata_field_info = [
    AttributeInfo(
        name="author",
        description="文档作者",
        type="string"
    ),
    AttributeInfo(
        name="date", 
        description="发布日期，格式: YYYY-MM-DD",
        type="string"
    ),
    AttributeInfo(
        name="category",
        description="文档类别，如: Python, 机器学习, Web开发等",
        type="string"
    ),
    AttributeInfo(
        name="views",
        description="浏览次数",
        type="integer"
    ),
    AttributeInfo(
        name="source",
        description="文档来源",
        type="string"
    ),
]

document_content_description = "技术文章、教程和研究论文"

# 创建自定义自查询检索器
self_query_retriever = CustomSelfQueryRetriever(
    llm=llm,
    vectorstore=vectorstore,  # 您的ChromaDB实例
    metadata_field_info=metadata_field_info,
    document_content_description=document_content_description
)

from langchain.chains.query_constructor.base import (
    StructuredQueryOutputParser,
    get_query_constructor_prompt,
)
from langchain_core.output_parsers import StrOutputParser
from vllm import SamplingParams
import re

class SimpleSelfQueryRetriever:
    """简化版自查询检索器"""
    
    def __init__(self, llm, vectorstore, metadata_field_info, document_content_description):
        self.llm = llm
        self.vectorstore = vectorstore
        self.metadata_field_info = metadata_field_info
        self.document_content_description = document_content_description
    
    def _call_llm(self, prompt: str, max_tokens: int = 300) -> str:
        """调用vLLM模型"""
        sampling_params = SamplingParams(
            temperature=0.1,
            top_p=0.9,
            max_tokens=max_tokens,
            stop=["<|im_end|>", "<|endoftext|>"]
        )
        
        outputs = self.llm.generate([prompt], sampling_params)
        if outputs and outputs[0].outputs:
            return outputs[0].outputs[0].text.strip()
        return ""
    
    def get_relevant_documents(self, query: str, k: int = 4):
        """简化版自查询"""
        print(f"🎯 查询: {query}")
        
        # 构建自查询提示词
        metadata_info = "\n".join([f"- {field.name}: {field.description} (类型: {field.type})" 
                                 for field in self.metadata_field_info])
        
        prompt = f"""<|im_start|>system
            你是一个查询分析助手。请分析用户查询并提取：
            
            1. 搜索关键词（用于语义搜索）
            2. 过滤条件（基于元数据）
            
            可用的元数据字段：
            {metadata_info}
            
            文档类型：{self.document_content_description}
            
            返回格式：
            关键词：[搜索关键词]
            过滤：[字段名 操作符 值] 或多个条件用逗号分隔
            
            示例：
            查询："阅读量超过100的技术文章"  
            返回：
            关键词：技术文章
            过滤：views > 100

            查询："语言为中文的技术文章"  
            返回：
            关键词：技术文章
            过滤：language == zh-CN
            
            只返回格式化的结果，不要解释。<|im_end|>
            <|im_start|>user
            查询：{query}<|im_end|>
            <|im_start|>assistant
            """
        
        try:
            # 调用LLM解析查询
            response = self._call_llm(prompt, max_tokens=200)
            print(f"🤖 LLM解析结果: {response}")
            
            # 解析响应
            search_keyword, filter_conditions = self._parse_llm_response(response, query)
            print(f"🔍 搜索关键词: {search_keyword}")
            print(f"📋 过滤条件: {filter_conditions}")
            
            # 构建过滤条件
            filter_dict = self._build_filter_dict(filter_conditions)
            
            # 执行检索
            if filter_dict:
                docs = self.vectorstore.similarity_search(
                    search_keyword, 
                    k=k, 
                    filter=filter_dict
                )
            else:
                docs = self.vectorstore.similarity_search(search_keyword, k=k)
            
            print(f"✅ 找到 {len(docs)} 个文档")
            return docs
            
        except Exception as e:
            print(f"❌ 自查询失败: {e}，使用普通检索")
            return self.vectorstore.similarity_search(query, k=k)
    
    def _parse_llm_response(self, response: str, original_query: str):
        """解析LLM响应"""
        # 默认值
        search_keyword = original_query
        filter_conditions = []
        
        lines = response.split('\n')
        for line in lines:
            if line.startswith('关键词：'):
                search_keyword = line.replace('关键词：', '').strip()
            elif line.startswith('过滤：'):
                filters = line.replace('过滤：', '').strip()
                if filters and filters != '无':
                    filter_conditions = [f.strip() for f in filters.split(',')]
        
        return search_keyword, filter_conditions
    
    def _build_filter_dict(self, filter_conditions):
        """构建过滤字典"""
        if not filter_conditions:
            return None
        
        filter_dict = {}
        
        for condition in filter_conditions:
            parts = condition.split()
            if len(parts) >= 3:
                field, op, value = parts[0], parts[1], ' '.join(parts[2:])
                
                # 处理值类型
                if value.isdigit():
                    value = int(value)
                
                op_map = {
                    '>': '$gt', '<': '$lt', '>=': '$gte', '<=': '$lte',
                    '=': '$eq', 'contains': '$regex'
                }
                
                if op in op_map:
                    if op == 'contains':
                        filter_dict[field] = {op_map[op]: f".*{value}.*"}
                    else:
                        filter_dict[field] = {op_map[op]: value}
        
        return filter_dict if filter_dict else None

# 简化版自查询检索器测试
def demo_simple_self_query():
    """演示简化版自查询检索器"""
    
    print("🚀 简化版自查询检索器测试")
    print("=" * 60)
    
    # 4. 定义元数据
    metadata_field_info = [
        AttributeInfo(
            name="author",
            description="文档作者",
            type="string"
        ),
        AttributeInfo(
            name="date", 
            description="发布日期，格式: YYYY-MM-DD",
            type="string"
        ),
        AttributeInfo(
            name="category",
            description="文档类别，如: Python, 机器学习, Web开发等",
            type="string" 
        ),
        AttributeInfo(
            name="views",
            description="浏览次数",
            type="integer"
        ),
        AttributeInfo(
            name="language",
            description="语言",
            type="string"
        ),
    ]
    
    document_content_description = "技术文章和教程"
    
    # 5. 直接创建简化版检索器
    print("📝 创建简化版自查询检索器...")
    retriever = SimpleSelfQueryRetriever(
        llm=llm,
        vectorstore=vectorstore, 
        metadata_field_info=metadata_field_info,
        document_content_description=document_content_description
    )
    
    print("✅ 简化版检索器创建成功！")
    
    # 6. 测试不同类型查询
    test_queries = [
        # 时间过滤查询
        {
            "query": "在CSDN上发布的博客",
            "description": "类别过滤"
        },
        # 数值过滤查询
        {
            "query": "收藏超过100的技术文章",
            "description": "数值范围过滤"
        },
        {
            "query": "语言为中文的技术文章",
            "description": "数值范围过滤"
        }
    ]
    
    print(f"\n🧪 开始测试 {len(test_queries)} 个查询...")
    
    for i, test_case in enumerate(test_queries, 1):
        query = test_case["query"]
        description = test_case["description"]
        
        print(f"\n" + "="*60)
        print(f"📋 测试 {i}/{len(test_queries)}: {description}")
        print(f"🔍 查询: {query}")
        
        # 执行查询
        results = retriever.get_relevant_documents(query, k=2)
        
        # 显示结果
        if results:
            print(f"✅ 找到 {len(results)} 个相关文档")
            for j, doc in enumerate(results, 1):
                # 清理输出
                preview = ' '.join(doc.page_content.split())[:80]
                print(f"   {j}. 标题: {preview}...")
                print(f"      元数据: {doc.metadata}")
        else:
            print("❌ 未找到相关文档")
    
    return retriever


# 运行测试
if __name__ == "__main__":
    # 测试简化版自查询检索器
    retriever = demo_simple_self_query()
    
    print("\n" + "="*60)
    print("🎉 简化版自查询检索器测试完成！")
    print("="*60)

🚀 简化版自查询检索器测试
============================================================
📝 创建简化版自查询检索器...
✅ 简化版检索器创建成功！

🧪 开始测试 3 个查询...

============================================================
📋 测试 1/3: 类别过滤
🔍 查询: 在CSDN上发布的博客
🎯 查询: 在CSDN上发布的博客
Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  2.94it/s]
🤖 LLM解析结果: 关键词：CSDN
             过滤：author == "CSDN"
🔍 搜索关键词: CSDN
📋 过滤条件: []
✅ 找到 2 个文档
✅ 找到 2 个相关文档
   1. 标题: 抓取检测之Dex-Net 2.0_dexnet-CSDN博客 抓取检测之Dex-Net 2.0 最新推荐文章于 2025-10-11 13:29:58 发布 原...
      元数据: {'title': '抓取检测之Dex-Net 2.0_dexnet-CSDN博客', 'language': 'zh-CN', 'source': 'https://blog.csdn.net/qq_40081208/article/details/111053208?ops_request_misc=&request_id=&biz_id=102&utm_term=%E6%8A%93%E5%8F%96%E6%A3%80%E6%B5%8B&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-1-111053208.142^v102^pc_search_result_base7&spm=1018.2226.3001.4187', 'description': '文章浏览阅读7.3k次，点赞13次，收藏72次。Dex-Net2.0是一种先进的机器人抓取算法，通过两阶段流程实现高效抓取：首先从深度图中采样抓取候选，接着评估抓取质量，最终选出最佳抓取配置。该算法引入了包含670万个样本的大规模数据集，并采用深度学习技术进行抓取质量评估。'}
   2. 标题: 机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks...
      元数据: {'description': '文章浏览阅读6.2k次，点赞10次，收藏97次。本文提出了一种基于深度学习的实时机器人抓取检测方法，利用卷积神经网络（CNN）进行单阶段回归预测，避免了传统滑动窗口方法的计算成本。模型在康奈尔抓取数据集上实现了88%的准确率，并能在GPU上以13帧/秒的速度运行。相比于现有技术，该模型提升了14个百分点的精度，并且能同时进行物体分类和抓取预测。此外，模型的多抓取版本可以预测单个物体的多个抓取点，显著提高了对多样抓取方式物体的检测性能。', 'title': '机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks-CSDN博客', 'language': 'zh-CN', 'source': 'https://blog.csdn.net/WhiffeYF/article/details/110829105'}

============================================================
📋 测试 2/3: 数值范围过滤
🔍 查询: 收藏超过100的技术文章
🎯 查询: 收藏超过100的技术文章
Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  3.47it/s]
🤖 LLM解析结果: 关键词：技术文章
             过滤：views > 100
🔍 搜索关键词: 技术文章
📋 过滤条件: []
✅ 找到 2 个文档
✅ 找到 2 个相关文档
   1. 标题: 机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks...
      元数据: {'language': 'zh-CN', 'title': '机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks-CSDN博客', 'description': '文章浏览阅读6.2k次，点赞10次，收藏97次。本文提出了一种基于深度学习的实时机器人抓取检测方法，利用卷积神经网络（CNN）进行单阶段回归预测，避免了传统滑动窗口方法的计算成本。模型在康奈尔抓取数据集上实现了88%的准确率，并能在GPU上以13帧/秒的速度运行。相比于现有技术，该模型提升了14个百分点的精度，并且能同时进行物体分类和抓取预测。此外，模型的多抓取版本可以预测单个物体的多个抓取点，显著提高了对多样抓取方式物体的检测性能。', 'source': 'https://blog.csdn.net/WhiffeYF/article/details/110829105'}
   2. 标题: 械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状 刘亚欣_基于深度图像的机械臂抓取位姿估计与轨迹优化研究-CSDN博客 械臂论文笔记（三）【抓取检测】...
      元数据: {'title': '械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状 刘亚欣_基于深度图像的机械臂抓取位姿估计与轨迹优化研究-CSDN博客', 'description': '文章浏览阅读9k次，点赞11次，收藏102次。本文综述了机器人抓取检测技术，重点介绍了基于学习的方法，包括基于抓取检测的抓取和基于视觉运动控制策略的端到端抓取，并探讨了各种方法的优势与局限。', 'language': 'zh-CN', 'source': 'https://blog.csdn.net/WhiffeYF/article/details/111031270?ops_request_misc=&request_id=&biz_id=102&utm_term=%E6%8A%93%E5%8F%96%E6%A3%80%E6%B5%8B&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-7-111031270.142^v102^pc_search_result_base7&spm=1018.2226.3001.4187'}

============================================================
📋 测试 3/3: 数值范围过滤
🔍 查询: 语言为中文的技术文章
🎯 查询: 语言为中文的技术文章
Processed prompts: 100%|██████████| 1/1 [00:00<00:00,  3.80it/s]
🤖 LLM解析结果: 关键词：技术文章
             过滤：language == zh-CN
🔍 搜索关键词: 技术文章
📋 过滤条件: []
✅ 找到 2 个文档
✅ 找到 2 个相关文档
   1. 标题: 机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks...
      元数据: {'description': '文章浏览阅读6.2k次，点赞10次，收藏97次。本文提出了一种基于深度学习的实时机器人抓取检测方法，利用卷积神经网络（CNN）进行单阶段回归预测，避免了传统滑动窗口方法的计算成本。模型在康奈尔抓取数据集上实现了88%的准确率，并能在GPU上以13帧/秒的速度运行。相比于现有技术，该模型提升了14个百分点的精度，并且能同时进行物体分类和抓取预测。此外，模型的多抓取版本可以预测单个物体的多个抓取点，显著提高了对多样抓取方式物体的检测性能。', 'title': '机械臂论文笔记（二）【实时抓取点检测】Real-Time Grasp Detection Using Convolutional Neural Networks-CSDN博客', 'language': 'zh-CN', 'source': 'https://blog.csdn.net/WhiffeYF/article/details/110829105'}
   2. 标题: 械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状 刘亚欣_基于深度图像的机械臂抓取位姿估计与轨迹优化研究-CSDN博客 械臂论文笔记（三）【抓取检测】...
      元数据: {'title': '械臂论文笔记（三）【抓取检测】机器人抓取检测技术的研究现状 刘亚欣_基于深度图像的机械臂抓取位姿估计与轨迹优化研究-CSDN博客', 'description': '文章浏览阅读9k次，点赞11次，收藏102次。本文综述了机器人抓取检测技术，重点介绍了基于学习的方法，包括基于抓取检测的抓取和基于视觉运动控制策略的端到端抓取，并探讨了各种方法的优势与局限。', 'source': 'https://blog.csdn.net/WhiffeYF/article/details/111031270?ops_request_misc=&request_id=&biz_id=102&utm_term=%E6%8A%93%E5%8F%96%E6%A3%80%E6%B5%8B&utm_medium=distribute.pc_search_result.none-task-blog-2~all~sobaiduweb~default-7-111031270.142^v102^pc_search_result_base7&spm=1018.2226.3001.4187', 'language': 'zh-CN'}

============================================================
🎉 简化版自查询检索器测试完成！
============================================================

from functools import lru_cache
import hashlib

class CachedRouter:
    """带缓存的路由器"""
    
    def __init__(self, semantic_router, cache_size=100):
        self.semantic_router = semantic_router
        self.cache_size = cache_size
        self._cache = {}
    
    def _get_cache_key(self, question: str) -> str:
        """生成缓存键"""
        return hashlib.md5(question.encode()).hexdigest()
    
    @lru_cache(maxsize=100)
    def route(self, question: str):
        """带缓存的路由"""
        cache_key = self._get_cache_key(question)
        
        if cache_key in self._cache:
            print("💾 使用缓存的路由结果")
            return self._cache[cache_key]
        
        # 执行路由
        route_name, score = self.semantic_router.route(question)
        
        # 存入缓存
        self._cache[cache_key] = (route_name, score)
        
        return route_name, score

import asyncio
from typing import List, Tuple

class ParallelRouter:
    """并行路由器"""
    
    async def route_multiple(
        self,
        questions: List[str]
    ) -> List[Tuple[str, str, float]]:
        """并行路由多个查询
        
        Returns:
            List[(question, route_name, score)]
        """
        async def route_single(question: str):
            # 异步路由（这里简化，实际需要异步嵌入）
            route_name, score = self.semantic_router.route(question)
            return question, route_name, score
        
        tasks = [route_single(q) for q in questions]
        results = await asyncio.gather(*tasks)
        
        return results

# 使用
# results = asyncio.run(router.route_multiple(["问题1", "问题2", "问题3"]))