从API到上下文：构建企业级RAG编程助手的完整技术实践

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import DirectoryLoader, TextLoader

# 加载并分割代码与文档
loader = DirectoryLoader("./repo/", glob="**/*.{md,py,js}", loader_cls=TextLoader)
docs = loader.load()

# 关键：针对代码语法定制分隔符
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=50,  # 重叠防止切断关键上下文
    separators=["\n\n", "\n", "def ", "class ", "```", " ", ""],
    keep_separator=True  # 保留def/class作为检索锚点
)
chunks = splitter.split_documents(docs)
print(f"Total chunks: {len(chunks)}")

from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma

# 加载本地BGE模型（无需联网，保障代码隐私）
embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-large-zh-v1.5",
    model_kwargs={'device': 'cuda'},
    encode_kwargs={'normalize_embeddings': True}
)

vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./db/chroma"
)
retriever = vectorstore.as_retriever(
    search_type="mmr",           # MMR减少冗余
    search_kwargs={"k": 6, "fetch_k": 20}
)

from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from langchain.chat_models import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

app = FastAPI()

def build_prompt(query: str, contexts: list) -> str:
    context_text = "\n\n---\n\n".join([doc.page_content for doc in contexts])
    return f"""【系统指令】你是严格的企业代码助手。请仅根据下方【参考上下文】回答。
如果上下文不足以回答问题，请明确回答“资料库中未找到相关信息”。
【参考上下文】
{context_text}
【用户问题】
{query}
【回答】"""

@app.post("/v1/chat/stream")
async def chat_stream(query: str):
    # 1. 检索增强
    docs = retriever.get_relevant_documents(query)
    
    # 2. 构建生成prompt
    full_prompt = build_prompt(query, docs)
    
    # 3. 初始化大模型（支持OpenAI或vLLM本地）
    llm = ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0.1,          # 低温度保证严谨性
        streaming=True,
        callbacks=[StreamingStdOutCallbackHandler()]
    )
    
    # 4. 流式响应
    async def generate():
        async for chunk in llm.astream(full_prompt):
            yield f"data: {chunk.content}\n\n"
        yield "data: [DONE]\n\n"
        
    return StreamingResponse(generate(), media_type="text/event-stream")

# 异步混合检索示例
async def hybrid_retrieve(query):
    vector_task = vectorstore.asimilarity_search(query, k=6)
    # 假设有关键词搜索引擎
    keyword_task = keyword_search(query, k=3) 
    results = await asyncio.gather(vector_task, keyword_task)
    return merge_and_deduplicate(results)