05_用LLM创建知识库：从文档到智能问答系统

用户提问 → 检索相关知识 → 结合知识与问题 → LLM生成回答

文档 → 文档处理 → 向量化 → 知识库存储
                                    ↑
                                    |
用户问题 → 查询处理 → 相似度检索 ←----------
                     ↓
              检索结果 + 问题 → LLM → 生成回答

import PyPDF2
import re
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 从PDF中提取文本
def extract_text_from_pdf(pdf_path):
    text = ""
    with open(pdf_path, 'rb') as file:
        reader = PyPDF2.PdfReader(file)
        for page_num in range(len(reader.pages)):
            page = reader.pages[page_num]
            text += page.extract_text()
    return text

# 清洗文本
def clean_text(text):
    # 去除多余的空白字符
    text = re.sub(r'\s+', ' ', text)
    # 可以根据需要添加更多清洗步骤
    return text

# 分段处理
def split_text(text, chunk_size=1000, chunk_overlap=200):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", " ", "", ""],
        length_function=len
    )
    chunks = text_splitter.split_text(text)
    return chunks

# 处理文档示例
pdf_path = "your_document.pdf"
raw_text = extract_text_from_pdf(pdf_path)
cleaned_text = clean_text(raw_text)
document_chunks = split_text(cleaned_text)

print(f"文档处理完成，共生成{len(document_chunks)}个文本片段")

from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np

# 加载嵌入模型和分词器
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-large-zh-v1.5")
model = AutoModel.from_pretrained("BAAI/bge-large-zh-v1.5")

# 确保模型在正确的设备上
if torch.cuda.is_available():
    model = model.to("cuda")

# 生成嵌入向量
def get_embeddings(texts, batch_size=32):
    embeddings = []
    # 批量处理文本以提高效率
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        # 分词
        inputs = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}
        
        # 获取模型输出
        with torch.no_grad():
            outputs = model(**inputs)
        
        # 计算句子嵌入（通常是[CLS]标记的最后一层隐藏状态）
        batch_embeddings = outputs.last_hidden_state[:, 0].cpu().numpy()
        
        # 归一化嵌入向量
        batch_embeddings = batch_embeddings / np.linalg.norm(batch_embeddings, axis=1, keepdims=True)
        embeddings.extend(batch_embeddings)
    
    return np.array(embeddings)

# 生成文档嵌入
document_embeddings = get_embeddings(document_chunks)
print(f"生成了{document_embeddings.shape[0]}个向量，每个向量维度为{document_embeddings.shape[1]}")

import chromadb
from chromadb.utils import embedding_functions

# 初始化Chroma DB客户端
client = chromadb.PersistentClient(path="./chroma_db")

# 创建或获取集合
collection_name = "my_knowledge_base"
collection = client.get_or_create_collection(
    name=collection_name,
    metadata={"description": "示例知识库"}
)

# 添加文档和嵌入到集合中
ids = [f"doc_{i}" for i in range(len(document_chunks))]
metadatas = [{"source": "your_document.pdf", "chunk_id": i} for i in range(len(document_chunks))]

collection.add(
    documents=document_chunks,
    embeddings=document_embeddings.tolist(),  # Chroma需要列表格式
    ids=ids,
    metadatas=metadatas
)

print(f"成功将{len(document_chunks)}个文档片段添加到知识库")

# 测试检索
query = "如何使用RAG技术？"
query_embedding = get_embeddings([query])[0]

results = collection.query(
    query_embeddings=[query_embedding.tolist()],
    n_results=3,
    include=["documents", "metadatas", "distances"]
)

print("检索结果：")
for i, (doc, meta, dist) in enumerate(zip(results["documents"][0], results["metadatas"][0], results["distances"][0])):
    print(f"\n结果 {i+1} (相似度: {1-dist:.4f}):")
    print(f"来源: {meta['source']}, 片段ID: {meta['chunk_id']}")
    print(f"内容: {doc[:150]}...")

def build_context(query, retrieved_docs, max_length=2000):
    """构建适合LLM的上下文"""
    # 按相似度排序（假设retrieved_docs已经排序）
    
    # 构建上下文
    context = ""
    context_parts = []
    
    for i, doc in enumerate(retrieved_docs):
        # 添加来源信息
        source_info = f"[来源: {doc['metadata'].get('source', 'Unknown')}, 片段: {doc['metadata'].get('chunk_id', i)}]\n"
        content = doc['document']
        
        # 检查添加当前文档是否会超出最大长度
        if len(context) + len(source_info) + len(content) + 10 > max_length:
            # 如果超出，尝试添加摘要
            if i == 0:  # 确保至少添加最相关的文档
                # 这里可以添加摘要生成逻辑
                content = content[:max_length - len(context) - len(source_info) - 3] + "..."
                context_parts.append(source_info + content)
            break
        
        context_parts.append(source_info + content)
    
    context = "\n\n---\n\n".join(context_parts)
    return context

# 示例使用
context = build_context(query, results_list)
print(f"构建的上下文长度: {len(context)} 字符")

def create_rag_prompt(query, context):
    prompt = f"""
你是一个专业的问答助手，你的任务是基于提供的上下文信息回答用户的问题。

请遵循以下规则：
1. 严格基于提供的上下文信息回答，不要添加上下文之外的信息
2. 如果上下文包含多个相关来源，请综合这些信息提供全面的回答
3. 如果上下文没有足够信息回答问题，请明确表示无法回答，不要编造信息
4. 在回答中引用相关信息来源

上下文信息：
{context}

用户问题：{query}

请基于上述上下文，提供详细、准确的回答。
"""
    return prompt

# 示例使用
rag_prompt = create_rag_prompt(query, context)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def generate_answer(prompt, model_name="meta-llama/Llama-3.1-70B-Instruct"):
    """使用LLM生成回答"""
    # 加载模型和分词器
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map="auto",
        load_in_4bit=True  # 使用4位量化以节省内存
    )
    
    # 准备输入
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # 生成回答
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=1000,
            temperature=0.3,  # 低温度以提高准确性
            top_p=0.9,
            use_cache=True
        )
    
    # 提取生成的回答
    answer = tokenizer.decode(
        outputs[0][inputs.input_ids.shape[1]:],
        skip_special_tokens=True
    )
    
    return answer

# 示例使用
answer = generate_answer(rag_prompt)
print("生成的回答:")
print(answer)

# 基础库
pip install python-dotenv pandas numpy matplotlib

# 文档处理
pip install PyPDF2 python-docx markdown openpyxl

# 文本处理和嵌入
pip install transformers langchain sentence-transformers nltk

# 向量数据库
pip install chromadb pymilvus-client

# LLM集成
pip install openai anthropic groq

import os
import re
import numpy as np
import PyPDF2
from transformers import AutoTokenizer, AutoModel
import torch
import chromadb
from langchain.text_splitter import RecursiveCharacterTextSplitter
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

class RAGSystem:
    def __init__(self, embedding_model="BAAI/bge-large-zh-v1.5", 
                 chroma_db_path="./chroma_knowledge_db",
                 collection_name="document_collection"):
        """初始化RAG系统"""
        self.embedding_model = embedding_model
        self.chroma_db_path = chroma_db_path
        self.collection_name = collection_name
        
        # 加载嵌入模型
        print("加载嵌入模型...")
        self.tokenizer = AutoTokenizer.from_pretrained(embedding_model)
        self.model = AutoModel.from_pretrained(embedding_model)
        if torch.cuda.is_available():
            self.model = self.model.to("cuda")
        
        # 初始化文本分块器
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            separators=["\n\n", "\n", " ", "", ""]
        )
        
        # 初始化向量数据库
        print("初始化向量数据库...")
        self.client = chromadb.PersistentClient(path=chroma_db_path)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"description": "RAG系统知识库"}
        )
        
        print("RAG系统初始化完成")
    
    def extract_text_from_pdf(self, pdf_path):
        """从PDF文件中提取文本"""
        text = ""
        try:
            with open(pdf_path, 'rb') as file:
                reader = PyPDF2.PdfReader(file)
                for page_num in range(len(reader.pages)):
                    page = reader.pages[page_num]
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n"
        except Exception as e:
            print(f"处理PDF文件时出错: {e}")
        return text
    
    def clean_text(self, text):
        """清洗文本"""
        # 去除多余的空白字符
        text = re.sub(r'\s+', ' ', text)
        # 去除特殊字符（可根据需要调整）
        text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
        return text.strip()
    
    def get_embeddings(self, texts, batch_size=32):
        """生成文本嵌入向量"""
        embeddings = []
        
        for i in range(0, len(texts), batch_size):
            batch = texts[i:i+batch_size]
            # 分词
            inputs = self.tokenizer(batch, padding=True, truncation=True, 
                                   max_length=512, return_tensors="pt")
            
            if torch.cuda.is_available():
                inputs = {k: v.to("cuda") for k, v in inputs.items()}
            
            # 获取模型输出
            with torch.no_grad():
                outputs = self.model(**inputs)
            
            # 使用[CLS]标记的嵌入作为句子嵌入
            batch_embeddings = outputs.last_hidden_state[:, 0].cpu().numpy()
            
            # 归一化嵌入向量
            batch_embeddings = batch_embeddings / np.linalg.norm(
                batch_embeddings, axis=1, keepdims=True
            )
            
            embeddings.extend(batch_embeddings)
        
        return np.array(embeddings)
    
    def add_document(self, file_path, metadata=None):
        """添加文档到知识库"""
        print(f"处理文档: {file_path}")
        
        # 根据文件类型提取文本
        file_ext = os.path.splitext(file_path)[1].lower()
        if file_ext == ".pdf":
            raw_text = self.extract_text_from_pdf(file_path)
        # 可以添加其他文件类型的处理逻辑
        else:
            print(f"不支持的文件类型: {file_ext}")
            return False
        
        # 清洗文本
        cleaned_text = self.clean_text(raw_text)
        
        # 分块处理
        chunks = self.text_splitter.split_text(cleaned_text)
        print(f"文档分割成 {len(chunks)} 个文本块")
        
        # 生成嵌入向量
        embeddings = self.get_embeddings(chunks)
        
        # 准备元数据
        if metadata is None:
            metadata = {}
        
        metadatas = []
        for i, chunk in enumerate(chunks):
            chunk_metadata = metadata.copy()
            chunk_metadata.update({
                "file_path": file_path,
                "chunk_id": i,
                "chunk_length": len(chunk),
                "chunk_start": chunk[:50] + "..." if len(chunk) > 50 else chunk
            })
            metadatas.append(chunk_metadata)
        
        # 生成文档ID
        base_id = os.path.splitext(os.path.basename(file_path))[0]
        ids = [f"{base_id}_chunk_{i}" for i in range(len(chunks))]
        
        # 添加到向量数据库
        self.collection.add(
            documents=chunks,
            embeddings=embeddings.tolist(),
            metadatas=metadatas,
            ids=ids
        )
        
        print(f"成功添加 {len(chunks)} 个文本块到知识库")
        return True
    
    def query_knowledge_base(self, query, n_results=5):
        """查询知识库"""
        # 生成查询嵌入
        query_embedding = self.get_embeddings([query])[0]
        
        # 查询向量数据库
        results = self.collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=n_results,
            include=["documents", "metadatas", "distances"]
        )
        
        # 整理结果
        formatted_results = []
        for i in range(len(results["documents"][0])):
            formatted_results.append({
                "document": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "distance": results["distances"][0][i],
                "similarity": 1 - results["distances"][0][i]
            })
        
        return formatted_results
    
    def build_context(self, retrieved_docs, max_length=2000):
        """构建上下文"""
        context_parts = []
        current_length = 0
        
        for i, doc in enumerate(retrieved_docs):
            # 添加来源信息
            source_info = f"[来源: {doc['metadata'].get('file_path', 'Unknown')}, "
            source_info += f"相关度: {doc['similarity']:.2f}]\n"
            
            content = doc['document']
            part_length = len(source_info) + len(content) + 10  # 加10作为分隔符
            
            # 检查是否会超出最大长度
            if current_length + part_length > max_length:
                if i == 0:  # 确保至少添加最相关的文档
                    # 截断内容以适应长度限制
                    available_space = max_length - current_length - len(source_info) - 3
                    if available_space > 0:
                        content = content[:available_space] + "..."
                        context_parts.append(source_info + content)
                        current_length += len(source_info) + len(content) + 10
                break
            
            context_parts.append(source_info + content)
            current_length += part_length
        
        context = "\n\n---\n\n".join(context_parts)
        return context
    
    def create_prompt(self, query, context):
        """创建提示模板"""
        prompt = f"""
你是一个专业的问答助手，负责基于提供的文档内容回答用户问题。

请严格按照以下要求回答：
1. 仅基于提供的上下文信息回答问题，不要添加上下文之外的信息
2. 如果你不确定答案或者上下文没有相关信息，请明确表示无法回答
3. 尽量提供详细、准确的回答，并在适当的地方引用文档来源
4. 保持回答的连贯性和可读性

上下文信息：
{context}

用户问题：{query}

请基于以上上下文，提供清晰、准确的回答。
"""
        return prompt
    
    def generate_answer(self, query, n_results=5, max_context_length=2000):
        """生成回答"""
        # 1. 查询知识库
        retrieved_docs = self.query_knowledge_base(query, n_results)
        
        # 2. 构建上下文
        context = self.build_context(retrieved_docs, max_context_length)
        
        # 3. 创建提示
        prompt = self.create_prompt(query, context)
        
        # 4. 调用LLM生成回答
        # 注意：这里需要根据实际使用的LLM进行调整
        # 以下是一个通用示例框架
        print("调用LLM生成回答...")
        # 示例：使用本地加载的模型
        # answer = self._generate_with_local_model(prompt)
        
        # 这里使用一个模拟的回答
        # 在实际应用中，应该替换为真实的LLM调用
        answer = "根据检索到的信息，这是对您问题的回答。在实际应用中，这里会是由LLM生成的完整回答。"
        
        # 5. 整理来源信息
        sources = []
        for doc in retrieved_docs:
            sources.append({
                "file_path": doc['metadata'].get('file_path', 'Unknown'),
                "similarity": doc['similarity'],
                "excerpt": doc['document'][:100] + "..." if len(doc['document']) > 100 else doc['document']
            })
        
        return {
            "query": query,
            "answer": answer,
            "sources": sources,
            "context": context
        }

# 使用示例
def main():
    # 创建RAG系统实例
    rag_system = RAGSystem()
    
    # 添加文档（示例文档路径）
    # rag_system.add_document("path/to/your/document.pdf")
    
    # 进行查询
    query = "什么是RAG技术？"
    result = rag_system.generate_answer(query)
    
    print("\n查询:")
    print(result["query"])
    print("\n回答:")
    print(result["answer"])
    print("\n来源:")
    for i, source in enumerate(result["sources"]):
        print(f"{i+1}. 文件: {source['file_path']}, 相关度: {source['similarity']:.2f}")
        print(f"   摘录: {source['excerpt']}")

if __name__ == "__main__":
    main()

def hybrid_search(self, query, n_results=5, keyword_weight=0.3):
    """混合向量检索和关键词检索"""
    # 向量检索
    vector_results = self.query_knowledge_base(query, n_results * 2)
    
    # 关键词检索（简化版，实际中可使用BM25等算法）
    keywords = query.split()
    keyword_scores = []
    for doc in vector_results:
        score = sum(1 for keyword in keywords if keyword.lower() in doc["document"].lower())
        keyword_scores.append((doc, score))
    
    # 混合排序
    hybrid_results = []
    for i, (doc, keyword_score) in enumerate(keyword_scores):
        # 结合向量相似度和关键词匹配分数
        hybrid_score = doc["similarity"] * (1 - keyword_weight) + \
                       (keyword_score / max(1, len(keywords))) * keyword_weight
        hybrid_results.append((doc, hybrid_score))
    
    # 按混合分数排序
    hybrid_results.sort(key=lambda x: x[1], reverse=True)
    
    return [doc for doc, score in hybrid_results[:n_results]]

def compress_context(self, context, target_length=1000):
    """压缩上下文，保留最重要的信息"""
    if len(context) <= target_length:
        return context
    
    # 简单的压缩策略：按段落重要性保留
    paragraphs = context.split("\n\n---\n\n")
    
    # 计算每个段落的重要性分数（示例：基于段落长度和位置）
    scores = []
    for i, para in enumerate(paragraphs):
        # 位置权重（前面的段落更重要）
        position_weight = 1.0 - (i * 0.1) if i < 10 else 0.1
        # 长度权重（更长的段落可能包含更多信息）
        length_weight = min(len(para) / 1000, 1.0)
        # 综合分数
        score = position_weight * 0.7 + length_weight * 0.3
        scores.append((para, score))
    
    # 按分数排序并选择最重要的段落
    scores.sort(key=lambda x: x[1], reverse=True)
    
    # 累积选择段落，直到达到目标长度
    compressed_context = []
    current_length = 0
    
    for para, score in scores:
        if current_length + len(para) + 10 <= target_length:  # +10 作为分隔符
            compressed_context.append(para)
            current_length += len(para) + 10
        elif compressed_context:  # 如果已经有内容，尝试截断当前段落
            remaining_space = target_length - current_length - 3
            if remaining_space > 50:  # 确保截断后仍有足够内容
                compressed_context.append(para[:remaining_space] + "...")
            break
        else:  # 如果还没有内容，必须添加至少一个段落
            compressed_context.append(para[:target_length-3] + "...")
            break
    
    return "\n\n---\n\n".join(compressed_context)

# 使用FastAPI部署RAG系统
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn

app = FastAPI(title="RAG智能问答系统")

# 全局RAG系统实例
rag_system = None

class QueryRequest(BaseModel):
    query: str
    n_results: int = 5
    max_context_length: int = 2000

class QueryResponse(BaseModel):
    query: str
    answer: str
    sources: list

@app.on_event("startup")
async def startup_event():
    global rag_system
    # 初始化RAG系统
    rag_system = RAGSystem()
    # 加载预设文档
    # rag_system.add_document("path/to/document.pdf")
    print("RAG系统已启动")

@app.post("/query", response_model=QueryResponse)
async def query_endpoint(request: QueryRequest):
    try:
        result = rag_system.generate_answer(
            query=request.query,
            n_results=request.n_results,
            max_context_length=request.max_context_length
        )
        return QueryResponse(
            query=result["query"],
            answer=result["answer"],
            sources=result["sources"]
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

class HierarchicalRAG:
    def __init__(self, base_rag):
        self.base_rag = base_rag  # 基础RAG系统
        self.document_level_collection = None
        self.paragraph_level_collection = None
        self.initialize_hierarchical_collections()
    
    def initialize_hierarchical_collections(self):
        """初始化分层集合"""
        # 文档级别集合
        self.document_level_collection = self.base_rag.client.get_or_create_collection(
            name="document_level_index",
            metadata={"description": "文档级索引"}
        )
        
        # 段落级别集合
        self.paragraph_level_collection = self.base_rag.client.get_or_create_collection(
            name="paragraph_level_index",
            metadata={"description": "段落级索引"}
        )
    
    def build_hierarchical_index(self, documents):
        """构建层次化索引"""
        for doc_path, doc_content in documents.items():
            # 1. 文档级索引
            doc_summary = self.generate_document_summary(doc_content)
            doc_embedding = self.base_rag.get_embeddings([doc_summary])[0]
            
            self.document_level_collection.add(
                documents=[doc_summary],
                embeddings=[doc_embedding.tolist()],
                ids=[f"doc_{os.path.basename(doc_path)}"],
                metadatas=[{"file_path": doc_path, "type": "document"}]
            )
            
            # 2. 段落级索引
            paragraphs = self.split_into_paragraphs(doc_content)
            for i, para in enumerate(paragraphs):
                para_embedding = self.base_rag.get_embeddings([para])[0]
                self.paragraph_level_collection.add(
                    documents=[para],
                    embeddings=[para_embedding.tolist()],
                    ids=[f"para_{os.path.basename(doc_path)}_{i}"],
                    metadatas=[{"file_path": doc_path, "paragraph_id": i, "type": "paragraph"}]
                )
    
    def hierarchical_search(self, query, top_docs=3, top_paras=5, top_chunks=10):
        """层次化搜索"""
        query_embedding = self.base_rag.get_embeddings([query])[0].tolist()
        
        # 1. 文档级搜索
        doc_results = self.document_level_collection.query(
            query_embeddings=[query_embedding],
            n_results=top_docs,
            include=["metadatas"]
        )
        
        # 2. 针对相关文档进行段落级搜索
        relevant_docs = [meta["file_path"] for meta in doc_results["metadatas"][0]]
        
        # 过滤获取相关文档的段落
        para_results = self.paragraph_level_collection.query(
            query_embeddings=[query_embedding],
            n_results=top_paras,
            where={"file_path": {"$in": relevant_docs}},
            include=["documents", "metadatas"]
        )
        
        # 3. 最后在基础RAG的细粒度块中搜索
        final_results = self.base_rag.query_knowledge_base(query, top_chunks)
        
        # 合并结果（可以根据需要实现更复杂的合并策略）
        # ...
        
        return final_results
    
    def generate_document_summary(self, content):
        """生成文档摘要"""
        # 这里可以实现文档摘要生成逻辑
        # ...
        return "文档摘要示例"
    
    def split_into_paragraphs(self, content):
        """将文本分割为段落"""
        # 简单实现，实际中可能需要更复杂的逻辑
        return [p.strip() for p in content.split("\n\n") if p.strip()]

class MultimodalRAG:
    def __init__(self):
        # 初始化文本嵌入模型
        self.text_encoder = ...  # 文本编码器
        
        # 初始化图像嵌入模型
        self.image_encoder = ...  # 图像编码器
        
        # 初始化向量数据库
        self.text_collection = ...  # 文本集合
        self.image_collection = ...  # 图像集合
        self.cross_media_collection = ...  # 跨媒体集合
    
    def add_image_document(self, image_path, caption=None):
        """添加图像到知识库"""
        # 提取图像特征
        image_features = self.extract_image_features(image_path)
        
        # 如果有标题，生成标题嵌入
        if caption:
            caption_embedding = self.text_encoder.encode([caption])[0]
        else:
            # 可以使用图像描述模型生成标题
            caption = self.generate_image_caption(image_path)
            caption_embedding = self.text_encoder.encode([caption])[0]
        
        # 计算多模态嵌入（可以是文本和图像嵌入的融合）
        multimodal_embedding = self.fuse_embeddings(caption_embedding, image_features)
        
        # 存储到数据库
        self.image_collection.add(
            ids=[image_path],
            embeddings=[image_features.tolist()],
            metadatas=[{"path": image_path, "caption": caption}]
        )
        
        self.cross_media_collection.add(
            ids=[f"mm_{image_path}"],
            embeddings=[multimodal_embedding.tolist()],
            metadatas=[{"type": "image", "path": image_path, "caption": caption}]
        )
    
    def multimodal_query(self, query, is_image_query=False, n_results=5):
        """多模态查询"""
        if is_image_query:
            # 图像处理逻辑
            query_embedding = self.extract_image_features(query)
        else:
            # 文本查询
            query_embedding = self.text_encoder.encode([query])[0]
        
        # 在跨媒体集合中查询
        results = self.cross_media_collection.query(
            query_embeddings=[query_embedding.tolist()],
            n_results=n_results,
            include=["metadatas"]
        )
        
        # 处理不同类型的结果
        processed_results = []
        for meta in results["metadatas"][0]:
            if meta["type"] == "image":
                # 获取图像相关信息
                # ...
                pass
            elif meta["type"] == "text":
                # 获取文本相关信息
                # ...
                pass
        
        return processed_results

class AdaptiveRAG:
    def __init__(self, base_rag):
        self.base_rag = base_rag
        self.query_classifier = self.initialize_query_classifier()
        self.strategy_registry = self.initialize_strategy_registry()
        self.feedback_memory = []
    
    def initialize_query_classifier(self):
        """初始化查询分类器"""
        # 可以是一个简单的规则系统，也可以是一个ML模型
        return SimpleQueryClassifier()
    
    def initialize_strategy_registry(self):
        """初始化策略注册表"""
        return {
            "fact_based": {
                "n_results": 3,
                "retrieval_type": "exact_match",
                "context_compression": "extractive",
                "temperature": 0.1
            },
            "complex_ reasoning": {
                "n_results": 7,
                "retrieval_type": "hybrid",
                "context_compression": "generative",
                "temperature": 0.3
            },
            "creative": {
                "n_results": 5,
                "retrieval_type": "semantic",
                "context_compression": "selective",
                "temperature": 0.7
            }
        }
    
    def analyze_query(self, query):
        """分析查询类型和特点"""
        query_type = self.query_classifier.classify(query)
        complexity = self.estimate_complexity(query)
        domain = self.detect_domain(query)
        
        return {
            "type": query_type,
            "complexity": complexity,
            "domain": domain
        }
    
    def select_strategy(self, query_analysis):
        """选择合适的检索和生成策略"""
        query_type = query_analysis["type"]
        
        # 基础策略
        strategy = self.strategy_registry.get(query_type, 
                                           self.strategy_registry["fact_based"]).copy()
        
        # 根据复杂度调整参数
        if query_analysis["complexity"] == "high":
            strategy["n_results"] = min(strategy["n_results"] * 2, 10)
            strategy["retrieval_type"] = "hybrid"
        
        # 根据领域进行特定调整
        # ...
        
        return strategy
    
    def generate_answer(self, query):
        """生成回答"""
        # 1. 分析查询
        analysis = self.analyze_query(query)
        
        # 2. 选择策略
        strategy = self.select_strategy(analysis)
        
        # 3. 根据策略执行检索
        if strategy["retrieval_type"] == "exact_match":
            docs = self.base_rag.query_knowledge_base(query, strategy["n_results"])
        elif strategy["retrieval_type"] == "hybrid":
            docs = self.base_rag.hybrid_search(query, strategy["n_results"])
        else:  # semantic
            docs = self.base_rag.query_knowledge_base(query, strategy["n_results"])
        
        # 4. 构建上下文（应用压缩策略）
        if strategy["context_compression"] == "extractive":
            context = self.base_rag.build_context(docs, 1500)
        elif strategy["context_compression"] == "generative":
            context = self.base_rag.build_context(docs, 2000)
            # 这里可以添加生成式压缩
        else:  # selective
            context = self.base_rag.build_context(docs, 1800)
        
        # 5. 生成回答（应用温度参数）
        prompt = self.base_rag.create_prompt(query, context)
        # 在实际生成中应用temperature参数
        answer = self.base_rag.generate_with_temperature(prompt, strategy["temperature"])
        
        return {
            "query": query,
            "answer": answer,
            "strategy_used": strategy,
            "query_analysis": analysis
        }
    
    def provide_feedback(self, query, answer, satisfaction_score, feedback_text=None):
        """收集用户反馈"""
        self.feedback_memory.append({
            "query": query,
            "answer": answer,
            "satisfaction": satisfaction_score,
            "feedback": feedback_text,
            "timestamp": time.time()
        })
        
        # 可以实现基于反馈的策略优化
        self.optimize_strategies_based_on_feedback()

class RAGAgent:
    def __init__(self, rag_system):
        self.rag_system = rag_system
        self.tools = self.initialize_tools()
        self.llm = self.initialize_llm()
        self.memory = []
    
    def initialize_tools(self):
        """初始化可用工具"""
        return {
            "search_knowledge_base": self.search_knowledge_base,
            "web_search": self.web_search,
            "calculate": self.calculate,
            # 可以添加更多工具
        }
    
    def initialize_llm(self):
        """初始化LLM"""
        # 返回LLM实例或接口
        return SomeLLMInterface()
    
    def search_knowledge_base(self, query, n_results=5):
        """搜索知识库工具"""
        return self.rag_system.query_knowledge_base(query, n_results)
    
    def web_search(self, query, n_results=3):
        """网络搜索工具"""
        # 实现网络搜索逻辑
        # ...
        return ["搜索结果1", "搜索结果2", "搜索结果3"]
    
    def calculate(self, expression):
        """计算工具"""
        # 安全的计算实现
        # ...
        return {"result": 42}
    
    def decide_next_action(self, query, context, history):
        """决定下一步行动"""
        decision_prompt = self.create_decision_prompt(query, context, history)
        decision = self.llm.generate(decision_prompt)
        
        # 解析决策结果
        # 例如：{"action": "tool_use", "tool": "search_knowledge_base", "params": {"query": "..."}}
        # 或：{"action": "direct_answer", "content": "..."}
        
        return self.parse_decision(decision)
    
    def execute(self, query, max_steps=5):
        """执行查询处理流程"""
        history = []
        context = ""
        
        for step in range(max_steps):
            # 决定下一步行动
            decision = self.decide_next_action(query, context, history)
            
            if decision["action"] == "direct_answer":
                # 直接生成回答
                return decision["content"]
            
            elif decision["action"] == "tool_use":
                # 使用工具
                tool_name = decision["tool"]
                tool_params = decision["params"]
                
                if tool_name in self.tools:
                    tool_result = self.tools[tool_name](**tool_params)
                    
                    # 更新上下文和历史
                    context += f"\n工具结果 ({tool_name}): {str(tool_result)}"
                    history.append({
                        "step": step,
                        "action": "tool_use",
                        "tool": tool_name,
                        "params": tool_params,
                        "result": tool_result
                    })
                else:
                    context += f"\n未知工具: {tool_name}"
            
        # 如果达到最大步骤，生成最终回答
        final_prompt = self.create_final_answer_prompt(query, context, history)
        return self.llm.generate(final_prompt)