AI大模型应用开发入门-LangChain实现RAG检索增强生成

# llm_env.py
from langchain.chat_models import init_chat_model
llm = init_chat_model("gpt-4o-mini", model_provider="openai") 

# rag.py
import os
import sys
import time
sys.path.append(os.getcwd())
from llm_set import llm_env
from langchain_openai import OpenAIEmbeddings
from langchain_postgres import PGVector
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langgraph.graph import START, StateGraph
from typing_extensions import List, TypedDict, Annotated
from typing import Literal
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, BaseMessage
from langchain_core.prompts import ChatPromptTemplate
# 初始化 LLM
llm = llm_env.llm
# 嵌入模型
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
# 向量数据库初始化
vector_store = PGVector(
    embeddings=embeddings,
    collection_name="my_rag_docs",
    connection="postgresql+psycopg2://postgres:123456@localhost:5433/langchainvector",
)
# 加载网页内容
url = "https://python.langchain.com/docs/tutorials/qa_chat_history/"
loader = WebBaseLoader(web_paths=(url,))
docs = loader.load()
for doc in docs:
    doc.metadata["source"] = url
# 文本分割
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50)
all_splits = text_splitter.split_documents(docs)
# 添加 section 元数据
total_documents = len(all_splits)
third = total_documents / 3
for i, document in enumerate(all_splits):
    if i < third:
        document.metadata["section"] = "beginning"
    elif i < 2 * third:
        document.metadata["section"] = "middle"
    else:
        document.metadata["section"] = "end"
# 检查是否已存在向量
existing = vector_store.similarity_search(url, k=1, filter={"source": url})
if not existing:
    _ = vector_store.add_documents(documents=all_splits)
    print("文档向量化完成") 

class Search(TypedDict):
    query: Annotated[str, "The question to be answered"]
    section: Annotated[
        Literal["beginning", "middle", "end"],
        ...,
        "Section to query.",
    ]
class State(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]
    query: Search
    context: List[Document]
    answer: set
# 分析意图 → 获取 query 与 section
def analyze(state: State):
    structtured_llm = llm.with_structured_output(Search)
    query = structtured_llm.invoke(state["messages"])
    return {"query": query}
# 相似度检索
def retrieve(state: State):
    query = state["query"]
    if hasattr(query, 'section'):
        filter = {"section": query["section"]}
    else:
        filter = None
    retrieved_docs = vector_store.similarity_search(query["query"], filter=filter)
    return {"context": retrieved_docs}
生成模块基于 
ChatPromptTemplate
 和当前上下文生成回答：
prompt_template = ChatPromptTemplate.from_messages(
    [
        ("system", "尽你所能按照上下文:{context}，回答问题：{question}。"),
    ]
)
def generate(state: State):
    docs_content = "\n\n".join(doc.page_content for doc in state["context"])
    messages = prompt_template.invoke({
        "question": state["query"]["query"],
        "context": docs_content,
    })
    response = llm.invoke(messages)
    return {"answer": response.content, "messages": [response]} 

graph_builder = StateGraph(State).add_sequence([analyze, retrieve, generate])
graph_builder.add_edge(START, "analyze") 

DB_URI = "postgresql://postgres:123456@localhost:5433/langchaindemo?sslmode=disable"
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup()
    graph = graph_builder.compile(checkpointer=checkpointer)
    input_thread_id = input("输入thread_id:")
    time_str = time.strftime("%Y%m%d", time.localtime())
    config = {"configurable": {"thread_id": f"rag-{time_str}-demo-{input_thread_id}"}}
    print("输入问题，输入 exit 退出。")
    while True:
        query = input("你: ")
        if query.strip().lower() == "exit":
            break
        input_messages = [HumanMessage(query)]
        response = graph.invoke({"messages": input_messages}, config=config)
        print(response["answer"])