312 lines
8.7 KiB
Python
312 lines
8.7 KiB
Python
|
"""
|
|||
|
|
支持 RAG 的多轮对话管理器。
|
|||
|
|
"""
|
|||
|
|
from typing import List, Dict, Any, Optional
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
|
|||
|
|
from app.db.models import Conversation, Message
|
|||
|
|
from app.services.model_manager import ModelManager
|
|||
|
|
from app.services.kb_manager import KnowledgeBaseManager
|
|||
|
|
from app.config import Settings
|
|||
|
|
from app.utils.exceptions import ResourceNotFoundError
|
|||
|
|
from app.utils.logger import get_logger
|
|||
|
|
|
|||
|
|
logger = get_logger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ConversationManager:
|
|||
|
|
"""
|
|||
|
|
管理具有可选 RAG 支持的多轮对话。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
db_session: Session,
|
|||
|
|
model_manager: ModelManager,
|
|||
|
|
kb_manager: KnowledgeBaseManager,
|
|||
|
|
settings: Settings,
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
初始化对话管理器。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
db_session: 数据库会话
|
|||
|
|
model_manager: 模型管理器实例
|
|||
|
|
kb_manager: 知识库管理器实例
|
|||
|
|
settings: 应用程序设置
|
|||
|
|
"""
|
|||
|
|
self.db = db_session
|
|||
|
|
self.model_manager = model_manager
|
|||
|
|
self.kb_manager = kb_manager
|
|||
|
|
self.settings = settings
|
|||
|
|
|
|||
|
|
def create_conversation(
|
|||
|
|
self, user_id: Optional[int] = None, title: str = "New Conversation"
|
|||
|
|
) -> Conversation:
|
|||
|
|
"""
|
|||
|
|
创建新对话。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
user_id: 可选的用户 ID
|
|||
|
|
title: 对话标题
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Conversation: 创建的对话
|
|||
|
|
"""
|
|||
|
|
logger.info("creating_conversation", user_id=user_id, title=title)
|
|||
|
|
|
|||
|
|
conversation = Conversation(user_id=user_id, title=title)
|
|||
|
|
self.db.add(conversation)
|
|||
|
|
self.db.commit()
|
|||
|
|
self.db.refresh(conversation)
|
|||
|
|
|
|||
|
|
logger.info("conversation_created", conversation_id=conversation.id)
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
def get_conversation(self, conversation_id: int) -> Conversation:
|
|||
|
|
"""
|
|||
|
|
根据 ID 获取对话。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Conversation: 对话实例
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ResourceNotFoundError: 如果找不到对话
|
|||
|
|
"""
|
|||
|
|
conversation = (
|
|||
|
|
self.db.query(Conversation)
|
|||
|
|
.filter(Conversation.id == conversation_id)
|
|||
|
|
.first()
|
|||
|
|
)
|
|||
|
|
if not conversation:
|
|||
|
|
raise ResourceNotFoundError(f"Conversation not found: {conversation_id}")
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
def add_message(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
role: str,
|
|||
|
|
content: str,
|
|||
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|||
|
|
) -> Message:
|
|||
|
|
"""
|
|||
|
|
向对话添加消息。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
role: 消息角色（'user'、'assistant'、'system'）
|
|||
|
|
content: 消息内容
|
|||
|
|
metadata: 可选的元数据
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Message: 创建的消息
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ResourceNotFoundError: 如果找不到对话
|
|||
|
|
"""
|
|||
|
|
# 验证对话是否存在
|
|||
|
|
self.get_conversation(conversation_id)
|
|||
|
|
|
|||
|
|
message = Message(
|
|||
|
|
conversation_id=conversation_id,
|
|||
|
|
role=role,
|
|||
|
|
content=content,
|
|||
|
|
msg_metadata=metadata or {},
|
|||
|
|
)
|
|||
|
|
self.db.add(message)
|
|||
|
|
self.db.commit()
|
|||
|
|
self.db.refresh(message)
|
|||
|
|
|
|||
|
|
logger.info(
|
|||
|
|
"message_added",
|
|||
|
|
conversation_id=conversation_id,
|
|||
|
|
role=role,
|
|||
|
|
message_id=message.id,
|
|||
|
|
)
|
|||
|
|
return message
|
|||
|
|
|
|||
|
|
def get_messages(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
limit: int = 50,
|
|||
|
|
offset: int = 0,
|
|||
|
|
) -> List[Message]:
|
|||
|
|
"""
|
|||
|
|
获取对话中的消息。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
limit: 最大消息数量
|
|||
|
|
offset: 分页偏移量
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
List[Message]: 消息列表
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ResourceNotFoundError: 如果找不到对话
|
|||
|
|
"""
|
|||
|
|
# 验证对话是否存在
|
|||
|
|
self.get_conversation(conversation_id)
|
|||
|
|
|
|||
|
|
messages = (
|
|||
|
|
self.db.query(Message)
|
|||
|
|
.filter(Message.conversation_id == conversation_id)
|
|||
|
|
.order_by(Message.created_at.asc())
|
|||
|
|
.offset(offset)
|
|||
|
|
.limit(limit)
|
|||
|
|
.all()
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return messages
|
|||
|
|
|
|||
|
|
async def chat(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
user_input: str,
|
|||
|
|
kb_id: Optional[int] = None,
|
|||
|
|
llm_name: Optional[str] = None,
|
|||
|
|
use_rag: bool = True,
|
|||
|
|
) -> Message:
|
|||
|
|
"""
|
|||
|
|
处理用户消息并生成助手回复。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
user_input: 用户消息
|
|||
|
|
kb_id: 可选的知识库 ID（用于 RAG）
|
|||
|
|
llm_name: 可选的 LLM 模型名称
|
|||
|
|
use_rag: 是否使用 RAG（默认 True）
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
Message: 助手的回复消息
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ResourceNotFoundError: 如果找不到对话
|
|||
|
|
"""
|
|||
|
|
logger.info(
|
|||
|
|
"processing_chat",
|
|||
|
|
conversation_id=conversation_id,
|
|||
|
|
use_rag=use_rag,
|
|||
|
|
kb_id=kb_id,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 验证对话是否存在
|
|||
|
|
conversation = self.get_conversation(conversation_id)
|
|||
|
|
|
|||
|
|
# 保存用户消息
|
|||
|
|
self.add_message(conversation_id, "user", user_input)
|
|||
|
|
|
|||
|
|
# 获取对话历史
|
|||
|
|
history = self._build_history(conversation_id, limit=10)
|
|||
|
|
|
|||
|
|
# 获取 LLM
|
|||
|
|
llm = self.model_manager.get_llm(llm_name)
|
|||
|
|
|
|||
|
|
# 如果启用了 RAG，从知识库构建上下文
|
|||
|
|
context = ""
|
|||
|
|
sources = []
|
|||
|
|
if use_rag and kb_id:
|
|||
|
|
try:
|
|||
|
|
results = self.kb_manager.query_kb(kb_id, user_input, k=3)
|
|||
|
|
context_parts = []
|
|||
|
|
for i, result in enumerate(results, 1):
|
|||
|
|
context_parts.append(
|
|||
|
|
f"[{i}] {result['content'][:500]}..."
|
|||
|
|
if len(result['content']) > 500
|
|||
|
|
else f"[{i}] {result['content']}"
|
|||
|
|
)
|
|||
|
|
sources.append(result['metadata'])
|
|||
|
|
context = "\n\n".join(context_parts)
|
|||
|
|
logger.info("rag_context_retrieved", kb_id=kb_id, source_count=len(results))
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning("rag_context_failed", error=str(e))
|
|||
|
|
context = "(No relevant context found)"
|
|||
|
|
|
|||
|
|
# 构建提示
|
|||
|
|
if context:
|
|||
|
|
prompt = f"""You are a helpful AI assistant. Use the following context from the knowledge base to answer the question.
|
|||
|
|
|
|||
|
|
Context:
|
|||
|
|
{context}
|
|||
|
|
|
|||
|
|
Conversation history:
|
|||
|
|
{history}
|
|||
|
|
|
|||
|
|
User: {user_input}
|
|||
|
|
Assistant:"""
|
|||
|
|
else:
|
|||
|
|
prompt = f"""{history}
|
|||
|
|
|
|||
|
|
User: {user_input}
|
|||
|
|
Assistant:"""
|
|||
|
|
|
|||
|
|
# 生成回复
|
|||
|
|
try:
|
|||
|
|
response = llm.predict(prompt)
|
|||
|
|
logger.info("llm_response_generated", conversation_id=conversation_id)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("llm_generation_failed", error=str(e))
|
|||
|
|
response = "I apologize, but I encountered an error generating a response. Please try again."
|
|||
|
|
|
|||
|
|
# 保存助手消息
|
|||
|
|
metadata = {}
|
|||
|
|
if sources:
|
|||
|
|
metadata["sources"] = sources
|
|||
|
|
if llm_name:
|
|||
|
|
metadata["model"] = llm_name
|
|||
|
|
|
|||
|
|
assistant_message = self.add_message(
|
|||
|
|
conversation_id, "assistant", response, metadata
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
logger.info("chat_complete", conversation_id=conversation_id)
|
|||
|
|
return assistant_message
|
|||
|
|
|
|||
|
|
def _build_history(self, conversation_id: int, limit: int = 10) -> str:
|
|||
|
|
"""
|
|||
|
|
构建对话历史字符串。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
limit: 要包含的最大消息数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
str: 格式化的对话历史
|
|||
|
|
"""
|
|||
|
|
messages = self.get_messages(conversation_id, limit=limit)
|
|||
|
|
|
|||
|
|
history_parts = []
|
|||
|
|
for msg in messages[:-1]: # 排除最后一条用户消息（已在提示中）
|
|||
|
|
if msg.role == "user":
|
|||
|
|
history_parts.append(f"User: {msg.content}")
|
|||
|
|
elif msg.role == "assistant":
|
|||
|
|
history_parts.append(f"Assistant: {msg.content}")
|
|||
|
|
|
|||
|
|
return "\n".join(history_parts) if history_parts else "(No previous messages)"
|
|||
|
|
|
|||
|
|
def delete_conversation(self, conversation_id: int) -> bool:
|
|||
|
|
"""
|
|||
|
|
删除对话及其所有消息。
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
conversation_id: 对话 ID
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
bool: 如果已删除则为 True
|
|||
|
|
|
|||
|
|
Raises:
|
|||
|
|
ResourceNotFoundError: 如果找不到对话
|
|||
|
|
"""
|
|||
|
|
conversation = self.get_conversation(conversation_id)
|
|||
|
|
|
|||
|
|
logger.info("deleting_conversation", conversation_id=conversation_id)
|
|||
|
|
|
|||
|
|
self.db.delete(conversation)
|
|||
|
|
self.db.commit()
|
|||
|
|
|
|||
|
|
logger.info("conversation_deleted", conversation_id=conversation_id)
|
|||
|
|
return True
|