chore: rename multiple files to improve importability and module structure
This commit is contained in:
Annyfee
@@ -1,64 +0,0 @@
|
||||
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}')
|
||||
@@ -1,78 +0,0 @@
|
||||
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("请先运行'00_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}')
|
||||
@@ -1,94 +0,0 @@
|
||||
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("请先运行'00_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}')
|
||||
@@ -1,108 +0,0 @@
|
||||
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}')
|
||||
@@ -1,170 +0,0 @@
|
||||
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']}")
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
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.
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user