实现LLM对话的记忆功能
用Flask框架和智谱AI的API,结合PostgreSQL数据库,构建一个功能完整的智能聊天应用。
大模型本身无法感知上下文,所以需要借助数据库来进行模拟人的记忆。
对话系统由三个核心模块组成:
-
Web服务层(Flask):处理HTTP请求和响应,提供RESTful API接口
-
AI模型层(ZhipuModel):与智谱AI API交互,处理自然语言理解和生成
-
数据存储层(MessageStorage):使用PostgreSQL数据库管理对话历史和会话状态
1数据库设计
和正常AI软件一样,侧边栏是会话列表,进入会话列表之后是对话,每次交互均需要记录AI和user的聊条,所以聊条需要一个自增ID。
CREATE TABLE sessions (
session_id VARCHAR(255) PRIMARY KEY, -- 会话唯一标识
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP -- 会话创建时间
);
CREATE TABLE chat (
id SERIAL PRIMARY KEY, -- 自增主键
session_id VARCHAR(255) REFERENCES sessions(session_id), -- 外键关联
role VARCHAR(10) NOT NULL, -- 角色: 'user' 或 'ai'
content TEXT NOT NULL, -- 消息内容
replyed_time TIME, -- 回复时间(只有时间)
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP -- 消息创建时间
);
通过历史对话的时候只需要发送session_id就可以检索到该会话下的所有聊条。



2实现过程
1. 用户请求 → Flask路由 → 参数验证
2. 生成/获取session_id → 存储用户消息到DB
3. 从DB获取历史记录 → 构建消息上下文
4. 调用ZhipuModel → 获取AI回复
5. 存储AI回复到DB → 返回响应给用户
当选中session_id的时候就表示调用先前已有的对话,就带有“记忆”功能了

3工作流程
storage.py
1. 消息存储功能 (store_message方法)
-
将聊天消息保存到数据库中
-
自动规范化角色名称(将"ai"或"assistant"转换为"ai",其他转换为"user")
-
自动处理会话记录,确保会话存在
-
使用事务保证数据一致性
2. 历史记录检索功能 (get_conversation_history方法)
-
根据会话ID获取聊天历史
-
支持限制返回的消息数量
-
按时间顺序(基于自增ID)返回消息
import psycopg2
from datetime import datetime, time
from typing import List, Dict, Any, Optional
class MessageStorage:
def __init__(self, host: str, dbname: str, user: str, password: str):
self.conn_config = {
'host': host,
'dbname': dbname,
'user': user,
'password': password,
'client_encoding': 'UTF-8'
}
def _get_connection(self):
return psycopg2.connect(**self.conn_config)
def store_message(
self,
session_id: str,
role: str,
content: str,
replyed_time: Optional[time] = None
) -> None:
normalized_role = 'ai' if role.lower() in ('ai', 'assistant') else 'user'
replyed_time = replyed_time or datetime.now().time()
with self._get_connection() as conn:
with conn.cursor() as cur:
try:
cur.execute("""
INSERT INTO session (session_id)
VALUES (%s) ON CONFLICT (session_id) DO NOTHING
""", (session_id,))
cur.execute("""
INSERT INTO chat (session_id, role, content, replyed_time)
VALUES (%s, %s, %s, %s)
""", (session_id, normalized_role, content, replyed_time))
conn.commit()
except Exception as e:
conn.rollback()
print(f"存储消息错误: {str(e)}")
raise
def get_conversation_history(
self,
session_id: str,
limit: int = 20
) -> List[Dict[str, Any]]:
with self._get_connection() as conn:
with conn.cursor() as cur:
cur.execute("""
SELECT session_id, role, content, replyed_time
FROM chat
WHERE session_id = %s
ORDER BY id ASC
LIMIT %s
""", (session_id, limit))
return [
{
'session_id': row[0],
'role': row[1],
'content': row[2],
'replyed_time': row[3]
}
for row in cur.fetchall()
]
volcengine_model.py
1. 大模型对话接口 (chat_completion方法)
-
提供统一的大模型调用接口
-
支持两种模式:流式和非流式
-
处理消息格式化和API调用
2. 非流式对话 (_normal_chat方法)
-
一次性获取完整的模型回复
-
返回包含内容和token用量的完整响应
-
适合需要完整回复后再处理的场景
3. 流式对话 (_stream_chat方法)
-
实时逐块返回模型生成的内容
-
使用生成器实现流式输出
-
适合需要实时显示生成过程的场景
import requests
import json
from typing import List, Dict, Any, Generator
class VolcengineModel:
def __init__(self, api_key: str, api_url: str = "https://ark.cn-beijing.volces.com/api/v3/chat/completions"):
self.api_key = api_key
self.api_url = api_url
self.headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
def chat_completion(self, messages: List[Dict[str, str]], stream: bool = False) -> Any:
"""与火山方舟大模型进行交互"""
payload = {
"model": "deepseek-v3-250324",
"messages": messages,
"temperature": 0,
"stream": stream
}
try:
if stream:
return self._stream_chat(payload)
else:
return self._normal_chat(payload)
except Exception as e:
raise Exception(f"API调用失败: {str(e)}")
def _normal_chat(self, payload: Dict[str, Any]) -> Dict[str, Any]:
"""非流式聊天"""
response = requests.post(
self.api_url,
headers=self.headers,
json=payload,
timeout=60
)
if response.status_code != 200:
raise Exception(f"API请求失败: {response.status_code} - {response.text}")
result = response.json()
return {
"content": result["choices"][0]["message"]["content"],
"usage": result.get("usage", {})
}
def _stream_chat(self, payload: Dict[str, Any]) -> Generator[Dict[str, Any], None, None]:
"""流式聊天"""
response = requests.post(
self.api_url,
headers=self.headers,
json=payload,
stream=True,
timeout=60
)
if response.status_code != 200:
raise Exception(f"API请求失败: {response.status_code} - {response.text}")
for line in response.iter_lines():
if line:
line = line.decode('utf-8')
if line.startswith('data: '):
data = line[6:] # 移除"data: "前缀
if data != '[DONE]':
chunk = json.loads(data)
if "choices" in chunk and len(chunk["choices"]) > 0:
delta = chunk["choices"][0].get("delta", {})
if "content" in delta:
yield {"content": delta["content"]}
app.py
1. POST /chat - 标准聊天接口
-
接收用户消息,返回完整的AI回复
-
自动管理会话ID
-
存储对话历史到数据库
2. POST /stream_chat - 流式聊天接口
-
使用Server-Sent Events (SSE) 实现实时流式输出
-
逐块返回AI生成的内容
-
支持事件驱动(session_info, error, end)
3. GET /history/<session_id> - 获取历史记录
-
查询指定会话的所有消息历史
-
自动处理字节编码问题
4. GET /test_db - 数据库连接测试
-
验证数据库连接状态
from flask import Flask, request, jsonify, Response, stream_with_context
from volcengine_model import VolcengineModel # 修改导入
from storage import MessageStorage
import uuid
import os
import json
import traceback
from typing import List, Dict
from datetime import datetime
app = Flask(__name__)
# 配置JSON处理确保UTF-8编码
app.json.ensure_ascii = False
app.json.mimetype = "application/json; charset=utf-8"
# 初始化模型和数据库 - 使用火山方舟API
volcengine_model = VolcengineModel(
api_key=os.getenv('VOLCENGINE_API_KEY', "9b6b8bd0- -4e53-ae63-91f89dbbddb8"),
api_url=os.getenv('VOLCENGINE_API_URL', "https://ark.cn-beijing.volces.com/api/v3/chat/completions")
)
# 初始化数据库
db_storage = MessageStorage(
host=os.getenv('DB_HOST', 'localhost'),
dbname=os.getenv('DB_NAME', 'db'),
user=os.getenv('DB_USER', 'postgres'),
password=os.getenv('DB_PASSWORD', '12345')
)
def prepare_messages(session_id: str, user_message: str) -> list:
"""准备对话历史消息"""
try:
history = db_storage.get_conversation_history(session_id)
messages = []
for item in history:
role = 'assistant' if item['role'] == 'ai' else 'user'
content = item['content'].decode('utf-8') if isinstance(item['content'], bytes) else item['content']
messages.append({
"role": role,
"content": content
})
messages.append({"role": "user", "content": user_message})
return messages
except Exception as e:
app.logger.error(f"准备消息历史出错: {str(e)}")
raise
@app.route('/chat', methods=['POST'])
def chat():
"""标准聊天接口"""
try:
data = request.get_json()
if not data or 'message' not in data:
return jsonify({"error": "缺少必要字段'message'"}), 400
session_id = data.get('session_id', str(uuid.uuid4()))
is_new_session = 'session_id' not in data
try:
db_storage.store_message(
session_id=session_id,
role='user',
content=data['message'],
replyed_time=datetime.now().time()
)
except Exception as e:
app.logger.error(f"存储用户消息失败: {str(e)}")
return jsonify({"error": "存储用户消息失败"}), 500
# 准备消息历史
messages = prepare_messages(session_id, data['message'])
# 获取AI回复 - 使用火山方舟模型
response = volcengine_model.chat_completion(messages=messages, stream=False)
try:
db_storage.store_message(
session_id=session_id,
role='ai',
content=response['content'],
replyed_time=datetime.now().time()
)
except Exception as e:
app.logger.error(f"存储AI回复失败: {str(e)}")
return jsonify({
"session_id": session_id,
"reply": response['content'],
"is_new_session": is_new_session
})
except Exception as e:
app.logger.error(f"聊天接口错误: {str(e)}\n{traceback.format_exc()}")
return jsonify({"error": "处理请求时出错"}), 500
@app.route('/stream_chat', methods=['POST'])
def stream_chat():
"""流式聊天接口"""
try:
data = request.get_json()
if not data or 'message' not in data:
return jsonify({"error": "缺少必要字段'message'"}), 400
session_id = data.get('session_id', str(uuid.uuid4()))
is_new_session = 'session_id' not in data
try:
db_storage.store_message(
session_id=session_id,
role='user',
content=data['message']
)
except Exception as e:
app.logger.error(f"存储用户消息失败: {str(e)}")
return jsonify({"error": "存储用户消息失败"}), 500
# 准备消息历史
messages = prepare_messages(session_id, data['message'])
def generate():
full_content = ""
try:
# 发送会话信息
yield f"event: session_info\ndata: {json.dumps({
'session_id': session_id,
'is_new_session': is_new_session
})}\n\n"
# 流式处理AI回复 - 使用火山方舟模型
for chunk in volcengine_model.chat_completion(messages=messages, stream=True):
content = chunk.get('content', '')
if content:
full_content += content
yield f"data: {json.dumps({
'text': content,
'session_id': session_id
})}\n\n"
try:
db_storage.store_message(
session_id=session_id,
role='ai',
content=full_content
)
except Exception as e:
app.logger.error(f"存储AI回复失败: {str(e)}")
yield "event: end\ndata: {\"status\": \"complete\"}\n\n"
except Exception as e:
error_msg = str(e)
yield f"event: error\ndata: {json.dumps({
'error': error_msg,
'session_id': session_id
})}\n\n"
app.logger.error(f"流式聊天错误: {error_msg}\n{traceback.format_exc()}")
return Response(
stream_with_context(generate()),
mimetype='text/event-stream',
headers={
'Cache-Control': 'no-cache',
'Connection': 'keep-alive'
}
)
except Exception as e:
app.logger.error(f"流式聊天初始化错误: {str(e)}\n{traceback.format_exc()}")
return jsonify({"error": "初始化流式聊天失败"}), 500
@app.route('/history/<session_id>', methods=['GET'])
def get_history(session_id):
"""获取会话历史"""
try:
history = db_storage.get_conversation_history(session_id)
# 编码清洗处理
def clean_data(item):
if isinstance(item, dict):
return {k: clean_data(v) for k, v in item.items()}
elif isinstance(item, bytes):
return item.decode('utf-8', errors='replace')
return item
safe_history = [clean_data(item) for item in history]
return jsonify({
"session_id": session_id,
"history": safe_history
})
except Exception as e:
app.logger.error(f"获取历史错误: {str(e)}\n{traceback.format_exc()}")
return jsonify({"error": "获取会话历史失败"}), 500
@app.route('/test_db', methods=['GET'])
def test_db():
"""测试数据库连接"""
try:
with db_storage._get_connection() as conn:
with conn.cursor() as cur:
cur.execute("SELECT 1")
result = cur.fetchone()
return jsonify({"status": "success", "result": result[0]})
except Exception as e:
return jsonify({"status": "error", "message": str(e)}), 500
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000, debug=True)
PS
7月份做的这个dome,当时zhipu api还有余额,现在没有了,改用了火山,所以调用的火山api还是显示zhipu的。

更多推荐


所有评论(0)