AI agent 大模型实战之 实战langchain集成RAG检索本地知识库

实战demo

import asyncio
import sys
from operator import itemgetter
from pathlib import Path
from langchain_core.prompts import ChatPromptTemplate, FewShotPromptTemplate, PromptTemplate, ChatMessagePromptTemplate
from langchain.agents import create_agent
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, TextLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

from common import llm

embedding_model = HuggingFaceEmbeddings(
    model_name=r'/Volumes/c/workspace/ai-agent-test/rag/model/BAAI/bge-large-zh-v1.5')

vector_store = Chroma(embedding_function=embedding_model, persist_directory='./chroma_v3')  # 数据库


def test():
    
    retriever = vector_store.as_retriever(search_kwargs={"k":5})

    prompt_template=PromptTemplate.from_template("""
    你是一个严谨的RAG助手
    请根据以下提供的上下文信息来回答问题
    如果上下文信息不足以回答问题，请直接说 无法回答
    如果回答使用了上下文中的信息，在回答后输出使用了哪些上下文
    上下文信息：
    {context}
    --------------------
    问题:{question}
    """)

    chain ={"question":RunnablePassthrough()} | RunnablePassthrough.assign(context=itemgetter("question") | retriever) | prompt_template | llm | StrOutputParser()

    resp = chain.invoke("周杰伦出生于什么地方")
    print(resp)


#初始化 知识库
async def init():
    # 获取本地知识库
    file_dir = Path("docs") 
    # 定义滑动窗口/切分
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)  
    # 加载文件
    docs = DirectoryLoader(str(file_dir),loader_cls=TextLoader).load()
    # 切分文件
    docs = text_splitter.split_documents(docs)
    # 存入数据库
    await vector_store.aadd_documents(docs)


#asyncio.run(init()) 执行一次知识库存进数据库
#init()
test()