构建可用于生产环境的 RAG 智能体：开发者完整指南

# 简单关键词搜索实现
def keyword_search(query, documents):
    query_terms = query.lower().split()
    scored_docs = []
    
    for doc in documents:
        score = sum(1 for term in query_terms if term in doc.lower())
        if score > 0:
            scored_docs.append((doc, score))
    
    return sorted(scored_docs, key=lambda x: x[1], reverse=True)

import chromadb
from chromadb.utils import embedding_functions

# Vector search 设置
client = chromadb.Client()
embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
    api_key="your-api-key",
    model_name="text-embedding-3-small"
)

collection = client.create_collection(
    name="knowledge_base",
    embedding_functinotallow=embedding_fn
)

defvector_search(query, top_k=5):
    results = collection.query(
        query_texts=[query],
        n_results=top_k
    )
    return results["documents"][0]

def reciprocal_rank_fusion(keyword_results, vector_results, k=60):
    """用 RRF 算法合并两个排名列表"""
    scores = {}
    
    # 给关键词结果打分
    for rank, doc inenumerate(keyword_results, 1):
        scores[doc] = scores.get(doc, 0) + 1 / (k + rank)
    
    # 给向量结果打分
    for rank, doc inenumerate(vector_results, 1):
        scores[doc] = scores.get(doc, 0) + 1 / (k + rank)
    
    # 按综合得分排序
    returnsorted(scores.items(), key=lambda x: x[1], reverse=True)

import chromadb
import openai
import tiktoken
from pathlib import Path
from typing importList

# 智能切块函数
defsmart_chunk(text: str, max_tokens: int = 200) -> List[str]:
    """按语义切分文本，限制 token 数"""
    tokenizer = tiktoken.get_encoding("cl100k_base")
    words = text.split()
    chunks = []
    current_chunk = []
    
    for word in words:
        # 检查加这个词会不会超 token 限制
        test_chunk = " ".join(current_chunk + [word])
        iflen(tokenizer.encode(test_chunk)) > max_tokens:
            if current_chunk:  # 不加空块
                chunks.append(" ".join(current_chunk))
            current_chunk = [word]
        else:
            current_chunk.append(word)
    
    # 加最后一块
    if current_chunk:
        chunks.append(" ".join(current_chunk))
    
    return chunks

# 加载并处理剧本
script_text = Path("back_to_the_future.txt").read_text(encoding="utf-8")
chunks = smart_chunk(script_text, max_tokens=200)
print(f"从剧本创建了 {len(chunks)} 个块")

from chromadb.utils import embedding_functions
import uuid

# 初始化带持久化的 ChromaDB
client = chromadb.Client(chromadb.config.Settings(
    persist_directory="./movie_knowledge_base"
))

# 设置 embedding 函数
embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
    api_key=openai.api_key,
    model_name="text-embedding-3-small"
)

# 创建或获取集合
collection_name = "bttf_script"
try:
    collection = client.get_collection(collection_name)
    print("找到已有集合")
except:
    collection = client.create_collection(
        name=collection_name,
        embedding_functinotallow=embedding_fn
    )
    print("创建新集合")

# 如果集合为空，添加文档
if collection.count() == 0:
    print("正在添加文档到集合...")
    collection.add(
        ids=[str(uuid.uuid4()) for _ in chunks],
        documents=chunks,
        metadatas=[{"chunk_index": i} for i inrange(len(chunks))]
    )
    client.persist()
    print("文档已添加并持久化")

def smart_search(query: str, top_k: int = 5) -> str:
    """
    多阶段检索：
    1. 用 vector search 广撒网
    2. 按关键词重叠 rerank
    3. 返回最佳结果
    """
    # 第一阶段：vector search 提高 recall
    initial_results = collection.query(
        query_texts=[query],
        n_results=min(top_k * 3, 15)  # 获取更多候选
    )
    
    docs = initial_results["documents"][0]
    ifnot docs:
        return"未找到相关信息。"
    
    # 第二阶段：按关键词重叠 rerank
    query_words = set(query.lower().split())
    scored_docs = []
    
    for doc in docs:
        doc_words = set(doc.lower().split())
        keyword_score = len(query_words.intersection(doc_words))
        scored_docs.append((doc, keyword_score))
    
    # 按关键词重叠排序，保留最佳结果
    reranked = sorted(scored_docs, key=lambda x: x[1], reverse=True)
    best_docs = [doc for doc, score in reranked[:top_k]]
    
    return"\n\n---\n\n".join(best_docs)

import uuid
from pathlib import Path
import chromadb
import tiktoken
from agents import Agent, Runner, function_tool
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 加载并切分剧本
script_text = Path("back_to_the_future.txt").read_text(encoding="utf-8")

defsimple_chunk(text, max_tokens=200):
    tokenizer = tiktoken.get_encoding("cl100k_base")
    words, chunk, chunks = text.split(), [], []
    for w in words:
        iflen(tokenizer.encode(" ".join(chunk + [w]))) > max_tokens:
            chunks.append(" ".join(chunk))
            chunk = [w]
        else:
            chunk.append(w)
    if chunk:
        chunks.append(" ".join(chunk))
    return chunks

docs = simple_chunk(script_text, max_tokens=200)

# 设置带持久化的 ChromaDB
client = chromadb.PersistentClient(path="./chroma_script_store")
collection_name = "bttf_script"

# 获取或创建集合
try:
    collection = client.get_collection(collection_name)
except Exception:
    collection = client.create_collection(name=collection_name)

# 如果集合为空，添加文档
if collection.count() == 0:
    collection.add(
        ids=[str(uuid.uuid4()) for _ in docs], 
        documents=docs
    )

# 定义搜索工具，带正确的装饰器
@function_tool
defsearch_script(query: str, top_k: int = 3) -> str:
    """搜索《回到未来》剧本中的相关片段"""
    res = collection.query(query_texts=[query], n_results=top_k)
    if res and"documents"in res and res["documents"] and res["documents"][0]:
        return"\n\n".join(res["documents"][0])
    return"未找到相关文档。"

# 创建智能体
agent = Agent(
    name="Script Agent",
    instructinotallow=(
        "你回答关于《回到未来》电影的问题。\n"
        "需要时调用 `search_script` 工具获取片段，"
        "然后引用或改述它们来回答。"
    ),
    tools=[search_script],
)

# 测试智能体
query = "Doc 在哪里让 Marty 见他，几点钟？"
result = Runner.run_sync(agent, query)
print("\n--- 答案 ---\n", result.final_output)

query = "凌晨 1:15 发生了什么？"
result = Runner.run_sync(agent, query)
print("\n--- 答案 ---\n", result.final_output)

# 生产级检索管道
def production_search(query: str):
    # 第一阶段：快速检索（广撒网）
    candidates = vector_search(query, k=20)
    
    # 第二阶段：关键词加权
    keyword_boosted = boost_keyword_matches(candidates, query)
    
    # 第三阶段：交叉编码器 rerank（如果预算够）
    final_results = cross_encoder_rerank(keyword_boosted, query, k=5)
    
    return final_results

def robust_search(query: str):
    try:
        return smart_search(query)
    except Exception as e:
        # 回退到简单搜索
        logging.error(f"智能搜索失败: {e}")
        return simple_keyword_search(query)