用Flask框架和智谱AI的API,结合PostgreSQL数据库,构建一个功能完整的智能聊天应用。

大模型本身无法感知上下文,所以需要借助数据库来进行模拟人的记忆。

对话系统由三个核心模块组成:

  1. Web服务层(Flask):处理HTTP请求和响应,提供RESTful API接口

  2. AI模型层(ZhipuModel):与智谱AI API交互,处理自然语言理解和生成

  3. 数据存储层(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的。

Logo

中国智能体开发者社区,聚焦智能体与大模型开发,提供前沿资讯、实用工具链、开源项目及行业案例。通过技术沙龙、开发者大赛等活动,促进经验交流与协作,助力开发者快速构建创新智能应用。

更多推荐