AI Agent(学习六-FAISS 持久化到磁盘(重启不丢记忆))
由于前段时间出差了半个月,今天才有空接着学,前面几篇我们已经具备的知识:
- 自动 Tool 选择
- 短期 Memory
- 长期 Memory(FAISS)
- 本地免费 Embedding
- 工程级 Prompt 拼装
目前的用例已经可以记住上下文的记忆,但是目前程序重启后会导致上一次的的记忆消失,所以,为了实现重启不丢长期记忆,需要将长期记忆存储到磁盘上,每次重启的时候从磁盘将数据加载到程序中
核心思路
- 程序启动时:如果磁盘上有FAISS,就load
- 每次写入长期记忆,立刻save到磁盘
实现
基于上一版本的代码,调整以下4步即可:
-
引入os
pythonimport os -
定义一个 FAISS 存储目录
在当前目录下创建一个存储目录
faiss_storepython#这个目录就是 Agent 的「长期大脑」 FAISS_PATH = "./faiss_store" -
改 FAISS 初始化
旧的代码
pythonvector_store = FAISS.from_texts( texts=["系统初始化占位文本"], embedding=embedding_model )修改的新代码
pythonif os.path.exists(FAISS_PATH): print("📦 加载已有长期记忆(FAISS)") vector_store = FAISS.load_local( FAISS_PATH, embedding_model, allow_dangerous_deserialization=True ) else: print("🆕 初始化新的长期记忆(FAISS)") vector_store = FAISS.from_texts( texts=["系统初始化占位文本"], embedding=embedding_model ) vector_store.save_local(FAISS_PATH)代码逻辑:
- load_local → 重启后恢复记忆
- allow_dangerous_deserialization=True是 LangChain 新版本的安全提示要求,本地用 OK
-
写入长期记忆时 save_local
旧的代码
pythondef save_long_term_memory(text: str): vector_store.add_texts([text])修改的新代码
pythondef save_long_term_memory(text: str): vector_store.add_texts([text]) vector_store.save_local(FAISS_PATH)代码解释:
每一次记忆写入,都会立刻保存
遇到问题
第一次启动时:
Traceback (most recent call last): File "/Users/wanping/IdeaProjects/ai/ai-agent/main5.py", line 25, in vector_store = FAISS.load_local( ^^^^^^^^^^^^^^^^^ File "/Users/wanping/IdeaProjects/ai/ai-agent/venv/lib/python3.12/site-packages/langchain_community/vectorstores/faiss.py", line 1205, in load_local index = faiss.read_index(str(path / f"{index_name}.faiss")) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/Users/wanping/IdeaProjects/ai/ai-agent/venv/lib/python3.12/site-packages/faiss/swigfaiss_avx2.py", line 13058, in read_index return _swigfaiss_avx2.read_index(*args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RuntimeError: Error in faiss::FileIOReader::FileIOReader(const char *) at /Users/runner/work/faiss-wheels/faiss-wheels/third-party/faiss/faiss/impl/io.cpp:70: Error: 'f' failed: could not open faiss_store/index.faiss for reading: No such file or directory
原因:代码逻辑是对的,FAISS 的"重启不丢"已经开始生效了,只是"第一次启动时磁盘上还没有索引文件"。
所以需要在加载磁盘文件的时候写try / except
注意:千万不要自己创建存储文件:index.faiss,index.pkl
可以修改为以下代码:
python
def load_or_create_faiss():
try:
if os.path.exists(FAISS_PATH):
print("📦 尝试加载 FAISS 长期记忆")
return FAISS.load_local(
FAISS_PATH,
embedding_model,
allow_dangerous_deserialization=True
)
else:
raise FileNotFoundError
except Exception as e:
print("⚠️ FAISS 加载失败,重建索引:", e)
shutil.rmtree(FAISS_PATH, ignore_errors=True)
vs = FAISS.from_texts(
["系统初始化占位文本"],
embedding_model
)
vs.save_local(FAISS_PATH)
return vs
vector_store = load_or_create_faiss()
运行与测试
运行
sh
python main5.py
测试
你:我叫wp
Agent: 你好,wp
你:我叫什么
Agent: 你好,wp
停了项目后重启项目
AI Agent(自动 Tool + Memory)已启动,输入 exit 退出
你:我叫什么
Agent: 你叫wp
完整版代码
python
import json
import os
import shutil
from langchain_groq import ChatGroq
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
load_dotenv()
# =========定义 FAISS 存储路径=======
FAISS_PATH = "./faiss_store"
# ========== 1️⃣ 初始化 LLM ==========
llm = ChatGroq(
model="llama-3.1-8b-instant",
temperature=0
)
# ===== 长期记忆(向量库) =====
embedding_model = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
)
def load_or_create_faiss():
try:
if os.path.exists(FAISS_PATH):
print("📦 尝试加载 FAISS 长期记忆")
return FAISS.load_local(
FAISS_PATH,
embedding_model,
allow_dangerous_deserialization=True
)
else:
raise FileNotFoundError
except Exception as e:
print("⚠️ FAISS 加载失败,重建索引:", e)
shutil.rmtree(FAISS_PATH, ignore_errors=True)
vs = FAISS.from_texts(
["系统初始化占位文本"],
embedding_model
)
vs.save_local(FAISS_PATH)
return vs
vector_store = load_or_create_faiss()
# ========== 2️⃣ 定义 Tools ==========
def calculator_tool(expression: str) -> str:
try:
result = eval(expression)
return f"计算结果:{result}"
except Exception as e:
return f"计算错误:{str(e)}"
def echo_tool(text: str) -> str:
return f"Echo Tool 收到:{text}"
def save_long_term_memory(text: str):
vector_store.add_texts([text])
vector_store.save_local(FAISS_PATH)
def search_long_term_memory(query: str, k: int = 3) -> str:
docs = vector_store.similarity_search(query, k=k)
if not docs:
return ""
result = "以下是你过往的重要信息:\n"
for doc in docs:
result += f"- {doc.page_content}\n"
return result
TOOLS = {
"calculator": calculator_tool,
"echo": echo_tool
}
# ========== 3️⃣ Memory(核心新增) ==========
conversation_memory = [] # List[{"role": "...", "content": "..."}]
MAX_MEMORY_LENGTH = 6 # 防止 prompt 无限变长
def add_to_memory(role: str, content: str):
conversation_memory.append({"role": role, "content": content})
# 只保留最近 N 条
if len(conversation_memory) > MAX_MEMORY_LENGTH:
del conversation_memory[0]
def render_memory() -> str:
"""
把 Memory 渲染成 prompt 文本
"""
if not conversation_memory:
return ""
memory_text = "以下是之前的对话内容(供你理解上下文):\n"
for msg in conversation_memory:
memory_text += f"{msg['role']}:{msg['content']}\n"
return memory_text
# ========== 4️⃣ Tool Router Prompt ==========
ROUTER_PROMPT = """
你是一个 AI Agent 的调度器。
你只能从以下工具中选择一个:
1. calculator:用于数学计算
2. echo:原样返回用户输入
请根据【上下文 + 用户输入】,判断是否需要使用工具。
如果需要,严格返回 JSON(不要解释):
{
"tool": "<工具名>",
"input": "<传给工具的参数>"
}
如果不需要任何工具,返回:
{
"tool": "none",
"input": "<直接回复用户的话>"
}
"""
# ========== 5️⃣ Agent 核心 ==========
def agent_run(user_input: str) -> str:
short_memory = render_memory()
long_memory = search_long_term_memory(user_input)
prompt = (
long_memory
+ short_memory
+ ROUTER_PROMPT
+ "用户输入:\n"
+ user_input
)
decision_msg = llm.invoke(prompt)
decision_text = decision_msg.content.strip()
try:
decision = json.loads(decision_text)
except json.JSONDecodeError:
reply_msg = llm.invoke(long_memory + short_memory + user_input)
reply = reply_msg.content
add_to_memory("用户", user_input)
add_to_memory("Agent", reply)
return reply
tool_name = decision.get("tool")
tool_input = decision.get("input", "")
if tool_name in TOOLS:
result = TOOLS[tool_name](tool_input)
add_to_memory("用户", user_input)
add_to_memory("Agent", result)
return result
reply = tool_input
# ⭐ 判断是否值得写入长期记忆(简单规则)
if "我叫" in user_input or "我是" in user_input:
save_long_term_memory(user_input)
add_to_memory("用户", user_input)
add_to_memory("Agent", reply)
return reply
# ========== 6️⃣ 启动 ==========
print("AI Agent(自动 Tool + Memory)已启动,输入 exit 退出")
while True:
user_input = input("你:")
if user_input in ("exit", "quit"):
print("Agent 已退出")
break
result = agent_run(user_input)
print("Agent:", result)