492 lines
16 KiB
Python
492 lines
16 KiB
Python
import json
|
||
import os
|
||
import uuid
|
||
from datetime import datetime
|
||
from typing import List, Dict, Optional, Any, Union
|
||
|
||
from peewee import (
|
||
Model, CharField, TextField, IntegerField,
|
||
DateTimeField, BooleanField, ForeignKeyField,
|
||
fn, SqliteDatabase
|
||
)
|
||
from playhouse.sqlite_ext import SqliteExtDatabase
|
||
from openai import OpenAI
|
||
|
||
|
||
# ==================== 数据库配置 ====================
|
||
class DatabaseManager:
|
||
"""数据库管理器"""
|
||
|
||
def __init__(self, db_path: str = 'chat_sessions.db'):
|
||
self.db = SqliteExtDatabase(
|
||
db_path,
|
||
pragmas={
|
||
'journal_mode': 'wal',
|
||
'cache_size': -1024 * 64,
|
||
'foreign_keys': 1,
|
||
'ignore_check_constraints': 0,
|
||
'synchronous': 1
|
||
}
|
||
)
|
||
self.init_models()
|
||
|
||
def init_models(self):
|
||
"""初始化数据模型"""
|
||
|
||
class BaseModel(Model):
|
||
class Meta:
|
||
database = self.db
|
||
|
||
class Session(BaseModel):
|
||
"""聊天会话模型"""
|
||
name = CharField(max_length=255, unique=True) # 会话名称唯一
|
||
model = CharField(max_length=50, default='deepseek-chat')
|
||
system_prompt = TextField(default='你是一个乐于助人的助手。')
|
||
created_at = DateTimeField(default=datetime.now)
|
||
updated_at = DateTimeField(default=datetime.now)
|
||
is_active = BooleanField(default=True)
|
||
metadata_json = TextField(default='{}')
|
||
|
||
@property
|
||
def metadata(self) -> Dict:
|
||
return json.loads(self.metadata_json) if self.metadata_json else {}
|
||
|
||
@metadata.setter
|
||
def metadata(self, value: Dict):
|
||
self.metadata_json = json.dumps(value, ensure_ascii=False)
|
||
|
||
def update_timestamp(self):
|
||
self.updated_at = datetime.now()
|
||
self.save()
|
||
|
||
class Message(BaseModel):
|
||
"""聊天消息模型"""
|
||
session = ForeignKeyField(Session, backref='messages', on_delete='CASCADE')
|
||
role = CharField(max_length=20, index=True)
|
||
content = TextField()
|
||
tokens = IntegerField(default=0)
|
||
timestamp = DateTimeField(default=datetime.now, index=True)
|
||
metadata_json = TextField(default='{}')
|
||
|
||
@property
|
||
def metadata(self) -> Dict:
|
||
return json.loads(self.metadata_json) if self.metadata_json else {}
|
||
|
||
@metadata.setter
|
||
def metadata(self, value: Dict):
|
||
self.metadata_json = json.dumps(value, ensure_ascii=False)
|
||
|
||
class Meta:
|
||
indexes = (
|
||
(('session', 'timestamp'), False),
|
||
)
|
||
|
||
# 保存模型类
|
||
self.Session = Session
|
||
self.Message = Message
|
||
|
||
# 创建表
|
||
self.db.connect()
|
||
self.db.create_tables([Session, Message], safe=True)
|
||
|
||
def close(self):
|
||
"""关闭数据库连接"""
|
||
if not self.db.is_closed():
|
||
self.db.close()
|
||
|
||
|
||
# ==================== 会话管理器 ====================
|
||
class SessionManager:
|
||
"""会话管理器"""
|
||
|
||
def __init__(self, db_path: str = None):
|
||
self.db_manager = DatabaseManager(db_path or 'chat_sessions.db')
|
||
self.Session = self.db_manager.Session
|
||
self.Message = self.db_manager.Message
|
||
|
||
def get_or_create_session(self, name: str, **kwargs) -> Any:
|
||
"""获取或创建会话"""
|
||
# 尝试获取现有会话
|
||
session = self.Session.get_or_none(self.Session.name == name)
|
||
|
||
if session:
|
||
return session
|
||
else:
|
||
# 创建新会话
|
||
defaults = {
|
||
'model': kwargs.get('model', 'deepseek-chat'),
|
||
'system_prompt': kwargs.get('system_prompt', '你是一个乐于助人的助手。'),
|
||
'metadata': kwargs.get('metadata', {})
|
||
}
|
||
|
||
session = self.Session.create(
|
||
name=name,
|
||
**defaults
|
||
)
|
||
|
||
# 添加系统提示作为第一条消息
|
||
if defaults['system_prompt']:
|
||
self.Message.create(
|
||
session=session,
|
||
role='system',
|
||
content=defaults['system_prompt']
|
||
)
|
||
|
||
return session
|
||
|
||
def add_message(self, session_name: str, role: str, content: str,
|
||
tokens: int = 0, metadata: Dict = None) -> str:
|
||
"""添加消息到会话"""
|
||
session = self.get_or_create_session(session_name)
|
||
|
||
self.Message.create(
|
||
session=session,
|
||
role=role,
|
||
content=content,
|
||
tokens=tokens,
|
||
timestamp=datetime.now(),
|
||
metadata=metadata or {}
|
||
)
|
||
|
||
# 更新会话时间戳
|
||
session.update_timestamp()
|
||
|
||
return content
|
||
|
||
def get_session_history(self, session_name: str, limit: int = 20) -> List[Dict]:
|
||
"""获取会话历史"""
|
||
session = self.Session.get_or_none(self.Session.name == session_name)
|
||
if not session:
|
||
return []
|
||
|
||
messages = (self.Message
|
||
.select()
|
||
.where(self.Message.session == session)
|
||
.order_by(self.Message.timestamp.asc())
|
||
.limit(limit))
|
||
|
||
history = []
|
||
for msg in messages:
|
||
history.append({
|
||
'role': msg.role,
|
||
'content': msg.content
|
||
})
|
||
|
||
return history
|
||
|
||
def get_session_info(self, session_name: str) -> Dict:
|
||
"""获取会话信息"""
|
||
session = self.Session.get_or_none(self.Session.name == session_name)
|
||
if not session:
|
||
return {}
|
||
|
||
messages = session.messages
|
||
message_count = messages.count()
|
||
user_messages = messages.where(self.Message.role == 'user').count()
|
||
assistant_messages = messages.where(self.Message.role == 'assistant').count()
|
||
|
||
return {
|
||
'name': session.name,
|
||
'model': session.model,
|
||
'system_prompt': session.system_prompt,
|
||
'message_count': message_count,
|
||
'user_messages': user_messages,
|
||
'assistant_messages': assistant_messages,
|
||
'created_at': session.created_at.isoformat(),
|
||
'updated_at': session.updated_at.isoformat(),
|
||
'metadata': session.metadata
|
||
}
|
||
|
||
def list_sessions(self) -> List[Dict]:
|
||
"""列出所有会话"""
|
||
sessions = self.Session.select().order_by(self.Session.updated_at.desc())
|
||
|
||
result = []
|
||
for session in sessions:
|
||
message_count = session.messages.count()
|
||
result.append({
|
||
'name': session.name,
|
||
'message_count': message_count,
|
||
'updated_at': session.updated_at.isoformat()
|
||
})
|
||
|
||
return result
|
||
|
||
def close(self):
|
||
"""关闭连接"""
|
||
self.db_manager.close()
|
||
|
||
|
||
# ==================== 简化的客户端接口 ====================
|
||
class DeepSeekChatClient:
|
||
"""简化的DeepSeek聊天客户端"""
|
||
|
||
def __init__(self, api_key: str, db_path: str = None):
|
||
"""
|
||
初始化客户端
|
||
|
||
Args:
|
||
api_key: DeepSeek API密钥
|
||
db_path: 数据库路径，默认为chat_sessions.db
|
||
"""
|
||
self.client = OpenAI(api_key=api_key, base_url='https://api.deepseek.com')
|
||
self.session_manager = SessionManager(db_path)
|
||
|
||
# 缓存已加载的会话配置
|
||
self._session_cache = {}
|
||
|
||
def chat(self, user_input: str, name: str, model: str = None,
|
||
system_prompt: str = None, **kwargs) -> str:
|
||
"""
|
||
发送聊天消息
|
||
|
||
Args:
|
||
user_input: 用户输入的问题
|
||
name: 会话名称
|
||
model: 模型名称（仅在创建新会话时使用）
|
||
system_prompt: 系统提示（仅在创建新会话时使用）
|
||
**kwargs: 其他OpenAI API参数
|
||
|
||
Returns:
|
||
AI的回复内容
|
||
"""
|
||
# 获取或创建会话
|
||
session = self.session_manager.get_or_create_session(
|
||
name=name,
|
||
model=model or 'deepseek-chat',
|
||
system_prompt=system_prompt
|
||
)
|
||
|
||
# 保存用户消息
|
||
self.session_manager.add_message(
|
||
session_name=name,
|
||
role='user',
|
||
content=user_input,
|
||
tokens=len(user_input) * 0.8
|
||
)
|
||
|
||
# 获取会话历史
|
||
history = self.session_manager.get_session_history(name, limit=10)
|
||
|
||
# 构建消息列表
|
||
messages = []
|
||
|
||
# 添加系统提示（如果历史中没有）
|
||
if not any(msg['role'] == 'system' for msg in history):
|
||
messages.append({
|
||
'role': 'system',
|
||
'content': session.system_prompt
|
||
})
|
||
|
||
# 添加上下文历史
|
||
for msg in history:
|
||
messages.append({
|
||
'role': msg['role'],
|
||
'content': msg['content']
|
||
})
|
||
|
||
# 添加当前用户消息
|
||
messages.append({
|
||
'role': 'user',
|
||
'content': user_input
|
||
})
|
||
|
||
# 调用API
|
||
try:
|
||
response = self.client.chat.completions.create(
|
||
model=session.model,
|
||
messages=messages,
|
||
max_tokens=kwargs.get('max_tokens', 2000),
|
||
temperature=kwargs.get('temperature', 0.7),
|
||
stream=kwargs.get('stream', False)
|
||
)
|
||
|
||
# 处理响应
|
||
ai_reply = response.choices[0].message.content
|
||
|
||
# 保存AI回复
|
||
self.session_manager.add_message(
|
||
session_name=name,
|
||
role='assistant',
|
||
content=ai_reply,
|
||
tokens=len(ai_reply) * 0.8
|
||
)
|
||
|
||
return ai_reply
|
||
|
||
except Exception as e:
|
||
error_msg = f"API调用错误: {str(e)}"
|
||
print(error_msg)
|
||
|
||
# 保存错误信息
|
||
self.session_manager.add_message(
|
||
session_name=name,
|
||
role='system',
|
||
content=error_msg,
|
||
metadata={'error': True, 'error_type': type(e).__name__}
|
||
)
|
||
|
||
return ""
|
||
|
||
def get_session_info(self, name: str) -> Dict:
|
||
"""获取会话信息"""
|
||
return self.session_manager.get_session_info(name)
|
||
|
||
def list_sessions(self) -> List[Dict]:
|
||
"""列出所有会话"""
|
||
return self.session_manager.list_sessions()
|
||
|
||
def export_session(self, name: str) -> Dict:
|
||
"""导出会话数据"""
|
||
info = self.get_session_info(name)
|
||
history = self.session_manager.get_session_history(name, limit=100)
|
||
|
||
return {
|
||
'session_info': info,
|
||
'history': history,
|
||
'exported_at': datetime.now().isoformat()
|
||
}
|
||
|
||
def clear_session_history(self, name: str, keep_system: bool = True) -> bool:
|
||
"""清空会话历史（保留系统提示）"""
|
||
session = self.session_manager.Session.get_or_none(
|
||
self.session_manager.Session.name == name
|
||
)
|
||
|
||
if not session:
|
||
return False
|
||
|
||
# 删除消息
|
||
query = session.messages
|
||
if keep_system:
|
||
query = query.where(self.session_manager.Message.role != 'system')
|
||
|
||
query.delete().execute()
|
||
|
||
# 重新添加系统提示（如果需要）
|
||
if keep_system and session.system_prompt:
|
||
messages = session.messages.where(
|
||
self.session_manager.Message.role == 'system'
|
||
)
|
||
if not messages.exists():
|
||
self.session_manager.add_message(
|
||
session_name=name,
|
||
role='system',
|
||
content=session.system_prompt
|
||
)
|
||
|
||
return True
|
||
|
||
def close(self):
|
||
"""关闭客户端"""
|
||
self.session_manager.close()
|
||
|
||
|
||
# ==================== 使用示例 ====================
|
||
def main():
|
||
# 初始化客户端（请使用环境变量管理API密钥）
|
||
API_KEY = "sk-bab97a2b9be042e18d945394f8feefa3"
|
||
|
||
client = None
|
||
try:
|
||
# 创建客户端
|
||
client = DeepSeekChatClient(api_key=API_KEY, db_path='my_chats.db')
|
||
print("✅ 客户端初始化成功！")
|
||
|
||
# ========== 第一次调用：创建会话 ==========
|
||
print("\n=== 第一次调用：创建技术讨论会话 ===")
|
||
res1 = client.chat(
|
||
name='技术讨论',
|
||
model='deepseek-chat',
|
||
system_prompt='你是一个技术专家，请提供详细的解释和代码示例。',
|
||
user_input='请解释Python中的装饰器。'
|
||
)
|
||
print("回答:", res1[:100] + "..." if len(res1) > 100 else res1)
|
||
|
||
# ========== 第二次调用：使用现有会话 ==========
|
||
print("\n=== 第二次调用：使用现有技术讨论会话 ===")
|
||
res2 = client.chat(
|
||
name='技术讨论', # 使用相同名称
|
||
user_input='请给我一个装饰器的实际应用例子。' # 只需要用户输入
|
||
)
|
||
print("回答:", res2[:100] + "..." if len(res2) > 100 else res2)
|
||
|
||
# ========== 第三次调用：创建新类型会话 ==========
|
||
print("\n=== 第三次调用：创建创意写作会话 ===")
|
||
res3 = client.chat(
|
||
name='创意写作',
|
||
model='deepseek-chat',
|
||
system_prompt='你是一个创意作家，请提供富有想象力的回答。',
|
||
user_input='写一个关于人工智能的短故事开头'
|
||
)
|
||
print("回答:", res3[:100] + "..." if len(res3) > 100 else res3)
|
||
|
||
# ========== 查看会话信息 ==========
|
||
print("\n=== 查看技术讨论会话信息 ===")
|
||
session_info = client.get_session_info('技术讨论')
|
||
print(f"会话名称: {session_info['name']}")
|
||
print(f"模型: {session_info['model']}")
|
||
print(f"消息总数: {session_info['message_count']}")
|
||
print(f"用户消息: {session_info['user_messages']}")
|
||
print(f"助手消息: {session_info['assistant_messages']}")
|
||
|
||
# ========== 列出所有会话 ==========
|
||
print("\n=== 所有会话列表 ===")
|
||
sessions = client.list_sessions()
|
||
for session in sessions:
|
||
print(f" {session['name']}: {session['message_count']} 条消息")
|
||
|
||
# ========== 导出会话数据 ==========
|
||
print("\n=== 导出技术讨论会话 ===")
|
||
export_data = client.export_session('技术讨论')
|
||
print(f"导出了 {len(export_data['history'])} 条消息")
|
||
|
||
print("\n🎉 测试完成！")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 程序运行出错: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
finally:
|
||
# 确保关闭连接
|
||
if client:
|
||
client.close()
|
||
|
||
|
||
# ==================== 简洁的使用示例 ====================
|
||
def simple_usage():
|
||
"""简洁的使用示例"""
|
||
API_KEY = "sk-bab97a2b9be042e18d945394f8feefa3"
|
||
|
||
# 1. 创建客户端
|
||
client = DeepSeekChatClient(api_key=API_KEY)
|
||
|
||
# 2. 创建新会话（第一次调用）
|
||
response1 = client.chat(
|
||
name='数学辅导',
|
||
model='deepseek-chat',
|
||
system_prompt='你是一个数学老师，请用简单易懂的方式解释数学概念。',
|
||
user_input='什么是微积分？'
|
||
)
|
||
print("第一次回答:", response1[:100])
|
||
|
||
# 3. 继续对话（只需要会话名称和用户输入）
|
||
response2 = client.chat(
|
||
name='数学辅导',
|
||
user_input='微积分有哪些实际应用？'
|
||
)
|
||
print("第二次回答:", response2[:100])
|
||
|
||
# 4. 查看会话信息
|
||
info = client.get_session_info('数学辅导')
|
||
print(f"会话 '{info['name']}' 有 {info['message_count']} 条消息")
|
||
|
||
client.close()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 运行完整示例
|
||
main()
|
||
|
||
# 或运行简洁示例
|
||
# simple_usage() |