chore: rename multiple files to improve importability and module structure
This commit is contained in:
Annyfee
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
from langchain_community.document_loaders import TextLoader
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from langchain_chroma import Chroma
|
||||
from embeddings import get_embeddings
|
||||
|
||||
|
||||
knowledge_base_file = "war_and_peace.txt"
|
||||
# 持久化目录: Chroma会把所有数据(向量+文本+元数据)都存到这个文件夹
|
||||
persist_directory = './chroma_db_war_and_peace_bge_small_en_v1.5'
|
||||
model_name_str = 'BAAI/bge-small-en-v1.5' # 如果愿意等待，可以换成模型"BAAI/bge-m3"，效果更好更适合长文，但下载时间也更久(2.2G)
|
||||
chunk_size = 500
|
||||
chunk_overlap = 75
|
||||
|
||||
|
||||
# 检查是否已创建
|
||||
if os.path.exists(persist_directory):
|
||||
print(f"检测到已存在的向量数据库: {persist_directory}")
|
||||
print("跳过索引构建。如需重新构建，请手动删除该目录。")
|
||||
exit()
|
||||
|
||||
if not os.path.exists(knowledge_base_file):
|
||||
print(f"错误: 知识库文件 {knowledge_base_file} 未找到。")
|
||||
print("请从 https://www.gutenberg.org/ebooks/2600.txt.utf-8 下载")
|
||||
print("并重命名为 war_and_peace.txt 放在当前目录。")
|
||||
exit()
|
||||
|
||||
print('---正在构建索引---')
|
||||
|
||||
# 1. 加载
|
||||
loader = TextLoader(knowledge_base_file,encoding='utf8')
|
||||
docs = loader.load()
|
||||
print('加载完成...\n')
|
||||
# 2. 分割
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap
|
||||
)
|
||||
splits = text_splitter.split_documents(docs)
|
||||
print('分割完成...\n')
|
||||
# 3. 向量化 -- 第一次运行会下载模型,预计耗时2分钟
|
||||
print(f'正在加载/下载模型{model_name_str}...')
|
||||
embedding_model = get_embeddings(
|
||||
model_name=model_name_str,
|
||||
device='cpu', # 强制模型在cpu上运行
|
||||
encode_kwargs={'batch_size':64} # 每次处理64个文本片段
|
||||
)
|
||||
print('Embedding模型加载完成...\n')
|
||||
# 4. 存储
|
||||
print('正在构建Chroma索引...(注：此步耗时较久，预计要3min)\n')
|
||||
db = Chroma(
|
||||
persist_directory=persist_directory,
|
||||
embedding_function=embedding_model
|
||||
)
|
||||
|
||||
# 分批添加切片chunks（每批不超过 5000）
|
||||
batch_size = 5000 # 必须 < 5461
|
||||
for i in range(0, len(splits), batch_size):
|
||||
batch = splits[i:i + batch_size]
|
||||
db.add_documents(batch)
|
||||
print(f"已插入 {min(i + batch_size, len(splits))} / {len(splits)} 条")
|
||||
|
||||
|
||||
print(f'✅ 索引构建完毕，共 {len(splits)} 条，已保存到 {persist_directory}')
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,78 @@
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_openai import ChatOpenAI
|
||||
from embeddings import get_embeddings
|
||||
from config import OPENAI_API_KEY
|
||||
import os
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
|
||||
|
||||
Persist_directory = './chroma_db_war_and_peace_bge_small_en_v1.5'
|
||||
model_name_str = 'BAAI/bge-small-en-v1.5'
|
||||
|
||||
if not os.path.exists(Persist_directory):
|
||||
print(f"错误: 知识库文件 {Persist_directory} 未找到。")
|
||||
print("请先运行'build_index.py'生成向量数据库，再运行该文件")
|
||||
exit()
|
||||
|
||||
print('---加载本地向量数据库---')
|
||||
|
||||
# 模块A:链接本地Chroma向量数据库
|
||||
# 1. 加载 Embedding 模型
|
||||
print(f'正在加载/下载模型{model_name_str}...')
|
||||
embeddings_model = get_embeddings(
|
||||
model_name=model_name_str,
|
||||
device='cpu'
|
||||
)
|
||||
|
||||
# 2. 从本地目录加载Chroma DB
|
||||
db = Chroma(
|
||||
persist_directory=Persist_directory,
|
||||
embedding_function=embeddings_model
|
||||
)
|
||||
print(f'Chroma数据库已从本地加载(共{db._collection.count()}条)\n')
|
||||
|
||||
# 模块B:R-A-G Flow
|
||||
# 1. R-检索
|
||||
retriever = db.as_retriever(search_kwargs={"k": 5}) # 召回5条相关数据
|
||||
|
||||
# 2. A-增强
|
||||
sys_prompt = """
|
||||
你是一个博学的历史学家和文学评论家。
|
||||
请根据以下上下文回答问题。如果上下文**强烈暗示**了答案，即使未明说，也可推理回答。
|
||||
如果完全无关，请回答“对不起，根据所提供的上下文我不知道”。
|
||||
|
||||
[上下文]: {context}
|
||||
[问题]: {question}
|
||||
"""
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
('system', sys_prompt),
|
||||
('human', '{question}')
|
||||
])
|
||||
|
||||
# 3. G-生成
|
||||
llm = ChatOpenAI(
|
||||
model="deepseek-chat",
|
||||
api_key=OPENAI_API_KEY,
|
||||
base_url="https://api.deepseek.com"
|
||||
)
|
||||
|
||||
# 4. 辅助函数
|
||||
def format_docs(docs):
|
||||
return "\n".join(doc.page_content for doc in docs)
|
||||
|
||||
|
||||
# 5. 组装RAG链条(LCEL)
|
||||
rag_chain = (
|
||||
{"context":retriever | format_docs, "question": RunnablePassthrough()}
|
||||
| prompt
|
||||
| llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
# 运行RAG链
|
||||
print('---正在运行RAG链条---')
|
||||
question = '莫斯科大火发生在小说的哪一部分？有哪些角色亲历了这场灾难？'
|
||||
response = rag_chain.invoke(question)
|
||||
print(f'提问:{question}')
|
||||
print(f'回答:{response}')
|
||||
@@ -0,0 +1,170 @@
|
||||
import os
|
||||
from config import OPENAI_API_KEY
|
||||
from embeddings import get_embeddings
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_classic.retrievers import ContextualCompressionRetriever
|
||||
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
|
||||
from langchain_classic.retrievers.document_compressors import CrossEncoderReranker
|
||||
from langchain_core.tools import tool
|
||||
from langchain_classic.agents import AgentExecutor
|
||||
from langchain_classic.agents import create_tool_calling_agent
|
||||
from langchain_community.chat_message_histories import ChatMessageHistory
|
||||
from langchain_core.runnables import RunnableWithMessageHistory
|
||||
|
||||
|
||||
|
||||
# 1. 构建一个可复用的 RAG链条
|
||||
def build_rag_chain(llm_instance):
|
||||
print('---正在构建RAG链条---')
|
||||
|
||||
persist_directory = './chroma_db_war_and_peace_bge_small_en_v1.5'
|
||||
embedding_model_name = 'BAAI/bge-small-en-v1.5'
|
||||
encoder_model_name = "BAAI/bge-reranker-base"
|
||||
|
||||
if not os.path.exists(persist_directory):
|
||||
raise FileNotFoundError(f'索引目录{persist_directory}未找到，请先运行 build_index.py')
|
||||
|
||||
# 链接向量数据库
|
||||
print(f'正在加载/下载 Embedding模型：{embedding_model_name}')
|
||||
embeddings_model = get_embeddings(model_name=embedding_model_name,device='cpu')
|
||||
db = Chroma(
|
||||
persist_directory=persist_directory,
|
||||
embedding_function=embeddings_model
|
||||
)
|
||||
# R
|
||||
print(f'正在加载 Reranker模型:{encoder_model_name}...')
|
||||
base_retriever = db.as_retriever(search_kwargs={'k':50})
|
||||
encoder = HuggingFaceCrossEncoder(model_name=encoder_model_name)
|
||||
reranker = CrossEncoderReranker(model=encoder,top_n=6)
|
||||
compression_retriever = ContextualCompressionRetriever(
|
||||
base_retriever=base_retriever,
|
||||
base_compressor=reranker
|
||||
)
|
||||
retriever = compression_retriever
|
||||
# A
|
||||
sys_prompt = """
|
||||
你是一个博学的历史学家和文学评论家。
|
||||
请根据以下上下文回答问题。如果上下文**强烈暗示**了答案，即使未明说，也可推理回答。
|
||||
如果完全无关，请回答“对不起，根据所提供的上下文我不知道”。
|
||||
|
||||
[上下文]: {context}
|
||||
[问题]: {question}
|
||||
"""
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
('system',sys_prompt),
|
||||
('human','{question}')
|
||||
])
|
||||
|
||||
def format_docs(docs):
|
||||
return '\n'.join(doc.page_content for doc in docs)
|
||||
|
||||
# R-A-G
|
||||
rag_chain = (
|
||||
{'context':retriever | format_docs,'question':RunnablePassthrough()}
|
||||
| prompt
|
||||
| llm_instance
|
||||
| StrOutputParser()
|
||||
)
|
||||
print('---RAG链条构建完毕!---\n')
|
||||
return rag_chain
|
||||
|
||||
|
||||
# 2. 将 RAG链条 组装进Agent里
|
||||
def create_agent_with_memory():
|
||||
# LLm
|
||||
llm = ChatOpenAI(
|
||||
model="deepseek-chat",
|
||||
api_key=OPENAI_API_KEY,
|
||||
base_url="https://api.deepseek.com"
|
||||
)
|
||||
# Prompt
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
('system','你是一个强大的助手。你能查天气，也能查《战争与和平》。请尽力回答用户所提的所有问题。'),
|
||||
MessagesPlaceholder(variable_name="history"), # 05篇所学:记忆占位符
|
||||
('human','{input}'),
|
||||
MessagesPlaceholder(variable_name="agent_scratchpad") # 05篇所学:ReAct 思考链，使其能够调用工具
|
||||
])
|
||||
|
||||
# Tool
|
||||
rag_chain_instance = build_rag_chain(llm_instance=llm)
|
||||
|
||||
|
||||
@tool
|
||||
def search_war_and_peace(query):
|
||||
"""查询《战争与和平》小说中的内容，包括人物、情节、历史事件等"""
|
||||
print(f'\n正在检索《战争与和平》:{query}')
|
||||
return rag_chain_instance.invoke(query)
|
||||
|
||||
@tool
|
||||
def get_weather(location):
|
||||
"""模拟获得天气信息"""
|
||||
return f"{location}当前天气：23℃，晴，风力2级"
|
||||
|
||||
tools = [get_weather,search_war_and_peace]
|
||||
|
||||
# 创建Agent
|
||||
agent = create_tool_calling_agent(llm=llm,tools=tools,prompt=prompt)
|
||||
agent_executor = AgentExecutor(agent=agent,tools=tools,verbose=False)
|
||||
|
||||
# 封装Memory
|
||||
store = {}
|
||||
|
||||
def get_session_history(session_id:int):
|
||||
if session_id not in store:
|
||||
store[session_id] = ChatMessageHistory()
|
||||
return store[session_id]
|
||||
|
||||
|
||||
# 添加记忆功能
|
||||
agent_with_memory = RunnableWithMessageHistory(
|
||||
runnable=agent_executor,
|
||||
get_session_history=get_session_history,
|
||||
input_messages_key="input",
|
||||
history_messages_key="history"
|
||||
)
|
||||
return agent_with_memory
|
||||
|
||||
|
||||
# 测试
|
||||
|
||||
if __name__ == '__main__':
|
||||
session_id = 'user123'
|
||||
agent = create_agent_with_memory()
|
||||
while 1:
|
||||
user_input = input('\n你:')
|
||||
if user_input=='quit':
|
||||
print('拜拜~')
|
||||
exit()
|
||||
response = agent.invoke(
|
||||
{'input':user_input},
|
||||
config={'configurable':{'session_id':session_id}}
|
||||
)
|
||||
print(f"AI:{response['output']}")
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
from config import OPENAI_API_KEY
|
||||
from embeddings import get_embeddings
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_classic.retrievers import ContextualCompressionRetriever
|
||||
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
|
||||
from langchain_classic.retrievers.document_compressors import CrossEncoderReranker
|
||||
from langchain_core.tools import tool
|
||||
|
||||
|
||||
# 全局 LLM (供Agent和Rag共用)
|
||||
llm = ChatOpenAI(
|
||||
model="deepseek-chat",
|
||||
api_key=OPENAI_API_KEY,
|
||||
base_url="https://api.deepseek.com"
|
||||
)
|
||||
|
||||
# (1) 构建一个可复用的 RAG链条 (P1+P2)
|
||||
def build_rag_chain(llm_instance):
|
||||
print('---正在构建RAG链条...---\n')
|
||||
|
||||
persist_directory = './chroma_db_war_and_peace_bge_small_en_v1.5'
|
||||
embedding_model_name = 'BAAI/bge-small-en-v1.5'
|
||||
encoder_model_name = "BAAI/bge-reranker-base"
|
||||
|
||||
if not os.path.exists(persist_directory):
|
||||
raise FileNotFoundError(f'索引目录{persist_directory}未找到，请先运行 build_index.py')
|
||||
|
||||
print(f'正在加载/下载 Embedding模型：{embedding_model_name}')
|
||||
embeddings_model = get_embeddings(model_name=embedding_model_name,device='cpu')
|
||||
db = Chroma(
|
||||
persist_directory=persist_directory,
|
||||
embedding_function=embeddings_model
|
||||
)
|
||||
|
||||
# 1. R-检索--强化版
|
||||
base_retriever = db.as_retriever(search_kwargs={"k":50})
|
||||
|
||||
print(f'正在加载 Reranker模型:{encoder_model_name}...')
|
||||
encoder = HuggingFaceCrossEncoder(model_name=encoder_model_name)
|
||||
reranker = CrossEncoderReranker(model=encoder,top_n=6)
|
||||
compression_retriever=ContextualCompressionRetriever(
|
||||
base_retriever=base_retriever,
|
||||
base_compressor=reranker
|
||||
)
|
||||
retriever = compression_retriever
|
||||
|
||||
# 2. A-增强
|
||||
sys_prompt = """
|
||||
你是一个博学的历史学家和文学评论家。
|
||||
请根据以下上下文回答问题。如果上下文**强烈暗示**了答案，即使未明说，也可推理回答。
|
||||
如果完全无关，请回答“对不起，根据所提供的上下文我不知道”。
|
||||
|
||||
[上下文]: {context}
|
||||
[问题]: {question}
|
||||
"""
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
('system',sys_prompt),
|
||||
('human','{question}')
|
||||
])
|
||||
# 3.G-生成(llm已在全局生成)
|
||||
|
||||
# 4. 辅助函数
|
||||
def format_docs(docs):
|
||||
return '\n'.join(doc.page_content for doc in docs)
|
||||
|
||||
# 5.组装RAG链条
|
||||
rag_chain = (
|
||||
{'context':retriever | format_docs, 'question': RunnablePassthrough()}
|
||||
| prompt
|
||||
| llm_instance
|
||||
| StrOutputParser()
|
||||
)
|
||||
print('---RAG链条构建完毕!---\n')
|
||||
return rag_chain
|
||||
|
||||
|
||||
# 初始化RAG链
|
||||
rag_chain_instance = build_rag_chain(llm)
|
||||
|
||||
# (2) 封装为标准 Langchain Tool
|
||||
@tool
|
||||
def search_war_and_peace(query):
|
||||
"""查询《战争与和平》小说中的内容，包括人物、情节、历史事件等"""
|
||||
print(f'\n正在检索《战争与和平》:{query}')
|
||||
return rag_chain_instance.invoke(query)
|
||||
|
||||
# 也可以与其他工具并列使用
|
||||
@tool
|
||||
def get_weather(location):
|
||||
"""模拟获得天气信息"""
|
||||
return f"{location}当前天气：23℃，晴，风力2级"
|
||||
|
||||
|
||||
tools = [search_war_and_peace,get_weather]
|
||||
|
||||
|
||||
# 运行
|
||||
if __name__ == '__main__':
|
||||
question = "皮埃尔是共济会成员吗？他在其中扮演什么角色？"
|
||||
res = search_war_and_peace.invoke(question)
|
||||
print(f'问题:{question}')
|
||||
print(f'回答:{res}')
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
from config import OPENAI_API_KEY
|
||||
from embeddings import get_embeddings
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_openai import ChatOpenAI
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.output_parsers import StrOutputParser
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_classic.retrievers import ContextualCompressionRetriever
|
||||
from langchain_classic.retrievers.document_compressors import CrossEncoderReranker
|
||||
from langchain_community.cross_encoders import HuggingFaceCrossEncoder
|
||||
|
||||
Persist_directory = './chroma_db_war_and_peace_bge_small_en_v1.5'
|
||||
model_name_str = 'BAAI/bge-small-en-v1.5'
|
||||
|
||||
if not os.path.exists(Persist_directory):
|
||||
print(f"错误: 知识库文件 {Persist_directory} 未找到。")
|
||||
print("请先运行'build_index.py'生成向量数据库，再运行该文件")
|
||||
exit()
|
||||
|
||||
print('---加载本地向量数据库---\n')
|
||||
|
||||
# 1. 加载 Embedding 模型
|
||||
print(f'正在加载/下载模型{model_name_str}...')
|
||||
embeddings_model = get_embeddings(
|
||||
model_name=model_name_str,
|
||||
device='cpu'
|
||||
)
|
||||
|
||||
# 2. 加载 Chroma db
|
||||
db = Chroma(
|
||||
persist_directory=Persist_directory,
|
||||
embedding_function=embeddings_model
|
||||
)
|
||||
|
||||
print('---Chroma数据库已加载---\n')
|
||||
|
||||
# --- 模块 B (R-A-G Flow) ---
|
||||
# 1. R-检索--强化版
|
||||
# 1.1 基础检索器(Base Retriever) - '粗召回'
|
||||
base_retriever = db.as_retriever(search_kwargs={"k":50}) # K调大到60
|
||||
# 1.2 Reranker (重排器) - "精排序" -- 首次运行需要耗时下载
|
||||
print('正在加载 Reranker模型 (bge-reranker-base)...')
|
||||
encoder = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base") # 加载Ranker模型
|
||||
reranker = CrossEncoderReranker(model=encoder,top_n=6) # 对检索结果进行精排
|
||||
# 1.3 创建管道封装器
|
||||
compression_retriever = ContextualCompressionRetriever(
|
||||
base_retriever=base_retriever, # 用Chroma做 海选
|
||||
base_compressor=reranker # 用Reranker做 精选
|
||||
)
|
||||
retriever = compression_retriever
|
||||
|
||||
print('--检索器已升级为Reranker模式--\n')
|
||||
|
||||
# 2. A-增强
|
||||
sys_prompt = """
|
||||
你是一个博学的历史学家和文学评论家。
|
||||
请根据以下上下文回答问题。如果上下文**强烈暗示**了答案，即使未明说，也可推理回答。
|
||||
如果完全无关，请回答“对不起，根据所提供的上下文我不知道”。
|
||||
|
||||
[上下文]: {context}
|
||||
[问题]: {question}
|
||||
"""
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
('system',sys_prompt),
|
||||
('human','{question}')
|
||||
])
|
||||
|
||||
# 3. G-生成
|
||||
llm = ChatOpenAI(
|
||||
model="deepseek-chat",
|
||||
api_key=OPENAI_API_KEY,
|
||||
base_url="https://api.deepseek.com"
|
||||
)
|
||||
|
||||
# 4. 辅助函数
|
||||
def format_docs(docs):
|
||||
return "\n".join(doc.page_content for doc in docs)
|
||||
|
||||
|
||||
# 5. 组装RAG链 (LCEL)
|
||||
rag_chain = (
|
||||
{"context":retriever | format_docs,"question":RunnablePassthrough()}
|
||||
| prompt
|
||||
| llm
|
||||
| StrOutputParser()
|
||||
)
|
||||
|
||||
# --运行RAG链--
|
||||
print('---正在运行RAG链条---')
|
||||
question = '皮埃尔是共济会成员吗？他在其中扮演什么角色？'
|
||||
response = rag_chain.invoke(question)
|
||||
print(f'提问:{question}')
|
||||
print(f'回答:{response}')
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user