295 lines
11 KiB
Python
295 lines
11 KiB
Python
|
"""Conversation service."""
|
|||
|
|
|
|||
|
|
from typing import List, Optional
|
|||
|
|
from sqlalchemy.orm import Session
|
|||
|
|
from sqlalchemy import select, desc, func, or_
|
|||
|
|
from langchain_core.messages import HumanMessage, AIMessage
|
|||
|
|
|
|||
|
|
from th_agenter.db.database import AsyncSessionFactory
|
|||
|
|
|
|||
|
|
from ..models.conversation import Conversation
|
|||
|
|
from ..models.message import Message, MessageRole
|
|||
|
|
from utils.util_schemas import ConversationCreate, ConversationUpdate
|
|||
|
|
from utils.util_exceptions import ConversationNotFoundError, DatabaseError
|
|||
|
|
from ..core.context import UserContext
|
|||
|
|
from datetime import datetime, timezone
|
|||
|
|
from loguru import logger
|
|||
|
|
|
|||
|
|
class ConversationService:
|
|||
|
|
"""Service for managing conversations and messages."""
|
|||
|
|
|
|||
|
|
def __init__(self, session: Session):
|
|||
|
|
self.session = session
|
|||
|
|
|
|||
|
|
async def create_conversation(
|
|||
|
|
self,
|
|||
|
|
user_id: int,
|
|||
|
|
conversation_data: ConversationCreate
|
|||
|
|
) -> Conversation:
|
|||
|
|
"""Create a new conversation."""
|
|||
|
|
self.session.desc = f"创建新会话 - 用户ID: {user_id}，会话数据: {conversation_data}"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
conversation = Conversation(
|
|||
|
|
**conversation_data.model_dump(),
|
|||
|
|
user_id=user_id
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Set audit fields
|
|||
|
|
conversation.set_audit_fields(user_id=user_id, is_update=False)
|
|||
|
|
|
|||
|
|
self.session.add(conversation)
|
|||
|
|
await self.session.commit()
|
|||
|
|
await self.session.refresh(conversation)
|
|||
|
|
|
|||
|
|
self.session.desc = f"创建新会话 Conversation ID: {conversation.id}，用户ID: {user_id}"
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
self.session.desc = f"ERROR: 创建会话失败 - 用户ID: {user_id}，错误: {str(e)}"
|
|||
|
|
await self.session.rollback()
|
|||
|
|
raise DatabaseError(f"创建会话失败: {str(e)}")
|
|||
|
|
|
|||
|
|
async def get_conversation(self, conversation_id: int) -> Optional[Conversation]:
|
|||
|
|
"""Get a conversation by ID."""
|
|||
|
|
try:
|
|||
|
|
user_id = UserContext.get_current_user_id()
|
|||
|
|
self.session.desc = f"获取会话 - 会话ID: {conversation_id}，用户ID: {user_id}"
|
|||
|
|
if user_id is None:
|
|||
|
|
logger.error(f"Failed to get conversation {conversation_id}: No user context available")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
conversation = await self.session.scalar(
|
|||
|
|
select(Conversation).where(
|
|||
|
|
Conversation.id == conversation_id,
|
|||
|
|
Conversation.user_id == user_id
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not conversation:
|
|||
|
|
self.session.desc = f"警告: 会话 {conversation_id} 不存在，用户ID: {user_id}"
|
|||
|
|
return conversation
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
self.session.desc = f"ERROR: 获取会话失败 - 会话ID: {conversation_id}，用户ID: {user_id}，错误: {str(e)}"
|
|||
|
|
raise DatabaseError(f"Failed to get conversation: {str(e)}")
|
|||
|
|
|
|||
|
|
async def get_user_conversations(
|
|||
|
|
self,
|
|||
|
|
skip: int = 0,
|
|||
|
|
limit: int = 50,
|
|||
|
|
search_query: Optional[str] = None,
|
|||
|
|
include_archived: bool = False,
|
|||
|
|
order_by: str = "updated_at",
|
|||
|
|
order_desc: bool = True
|
|||
|
|
) -> List[Conversation]:
|
|||
|
|
"""Get user's conversations with search and filtering."""
|
|||
|
|
user_id = UserContext.get_current_user_id()
|
|||
|
|
if user_id is None:
|
|||
|
|
logger.error("Failed to get user conversations: No user context available")
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
query = select(Conversation).where(
|
|||
|
|
Conversation.user_id == user_id
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Filter archived conversations
|
|||
|
|
if not include_archived:
|
|||
|
|
query = query.where(Conversation.is_archived == False)
|
|||
|
|
|
|||
|
|
# Search functionality
|
|||
|
|
if search_query and search_query.strip():
|
|||
|
|
search_term = f"%{search_query.strip()}%"
|
|||
|
|
query = query.where(
|
|||
|
|
or_(
|
|||
|
|
Conversation.title.ilike(search_term),
|
|||
|
|
Conversation.system_prompt.ilike(search_term)
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Ordering
|
|||
|
|
order_column = getattr(Conversation, order_by, Conversation.updated_at)
|
|||
|
|
if order_desc:
|
|||
|
|
query = query.order_by(desc(order_column))
|
|||
|
|
else:
|
|||
|
|
query = query.order_by(order_column)
|
|||
|
|
|
|||
|
|
return (await self.session.scalars(query.offset(skip).limit(limit))).all()
|
|||
|
|
|
|||
|
|
async def update_conversation(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
conversation_update: ConversationUpdate
|
|||
|
|
) -> Optional[Conversation]:
|
|||
|
|
"""Update a conversation."""
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if not conversation:
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
update_data = conversation_update.dict(exclude_unset=True)
|
|||
|
|
for field, value in update_data.items():
|
|||
|
|
setattr(conversation, field, value)
|
|||
|
|
|
|||
|
|
# Update audit fields
|
|||
|
|
conversation.set_audit_fields(user_id=conversation.user_id, is_update=True)
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
await self.session.commit()
|
|||
|
|
await self.session.refresh(conversation)
|
|||
|
|
return conversation
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to update conversation {conversation_id}: {str(e)}", exc_info=True)
|
|||
|
|
await self.session.rollback()
|
|||
|
|
raise DatabaseError(f"Failed to update conversation: {str(e)}")
|
|||
|
|
|
|||
|
|
async def delete_conversation(self, conversation_id: int) -> bool:
|
|||
|
|
"""Delete a conversation."""
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if not conversation:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
await self.session.delete(conversation)
|
|||
|
|
await self.session.commit()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def get_conversation_messages(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
skip: int = 0,
|
|||
|
|
limit: int = 100
|
|||
|
|
) -> List[Message]:
|
|||
|
|
"""Get messages from a conversation."""
|
|||
|
|
return (await self.session.scalars(
|
|||
|
|
select(Message).where(
|
|||
|
|
Message.conversation_id == conversation_id
|
|||
|
|
).order_by(Message.created_at).offset(skip).limit(limit)
|
|||
|
|
)).all()
|
|||
|
|
|
|||
|
|
async def add_message(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
content: str,
|
|||
|
|
role: MessageRole,
|
|||
|
|
message_metadata: Optional[dict] = None,
|
|||
|
|
context_documents: Optional[list] = None,
|
|||
|
|
prompt_tokens: Optional[int] = None,
|
|||
|
|
completion_tokens: Optional[int] = None,
|
|||
|
|
total_tokens: Optional[int] = None
|
|||
|
|
) -> Message:
|
|||
|
|
"""Add a message to a conversation."""
|
|||
|
|
message = Message(
|
|||
|
|
conversation_id=conversation_id,
|
|||
|
|
content=content,
|
|||
|
|
role=role,
|
|||
|
|
message_metadata=message_metadata,
|
|||
|
|
context_documents=context_documents,
|
|||
|
|
prompt_tokens=prompt_tokens,
|
|||
|
|
completion_tokens=completion_tokens,
|
|||
|
|
total_tokens=total_tokens
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Set audit fields
|
|||
|
|
message.set_audit_fields()
|
|||
|
|
|
|||
|
|
session = AsyncSessionFactory()
|
|||
|
|
session.begin()
|
|||
|
|
try:
|
|||
|
|
session.add(message)
|
|||
|
|
await session.commit()
|
|||
|
|
await session.refresh(message)
|
|||
|
|
|
|||
|
|
# Update conversation's updated_at timestamp
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if conversation:
|
|||
|
|
conversation.updated_at = datetime.now(timezone.utc)
|
|||
|
|
conversation.set_audit_fields(user_id=conversation.user_id, is_update=True)
|
|||
|
|
await session.commit()
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to add message to conversation {conversation_id}: {str(e)}", exc_info=True)
|
|||
|
|
await session.rollback()
|
|||
|
|
finally:
|
|||
|
|
await session.close()
|
|||
|
|
|
|||
|
|
return message
|
|||
|
|
|
|||
|
|
async def get_conversation_history_messages(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
limit: int = 20
|
|||
|
|
) -> List[Message]:
|
|||
|
|
"""Get recent conversation history messages."""
|
|||
|
|
history = await self.get_conversation_history(conversation_id, limit)
|
|||
|
|
history_messages = []
|
|||
|
|
for message in history:
|
|||
|
|
if message.role == MessageRole.USER:
|
|||
|
|
history_messages.append(HumanMessage(content=message.content))
|
|||
|
|
elif message.role == MessageRole.ASSISTANT:
|
|||
|
|
history_messages.append(AIMessage(content=message.content))
|
|||
|
|
return history_messages
|
|||
|
|
|
|||
|
|
async def get_conversation_history(
|
|||
|
|
self,
|
|||
|
|
conversation_id: int,
|
|||
|
|
limit: int = 20
|
|||
|
|
) -> List[Message]:
|
|||
|
|
"""Get recent conversation history for context."""
|
|||
|
|
return (await self.session.scalars(
|
|||
|
|
select(Message).where(
|
|||
|
|
Message.conversation_id == conversation_id
|
|||
|
|
).order_by(desc(Message.created_at)).limit(limit)
|
|||
|
|
)).all()[::-1] # Reverse to get chronological order
|
|||
|
|
|
|||
|
|
async def update_conversation_timestamp(self, conversation_id: int) -> None:
|
|||
|
|
"""Update conversation's updated_at timestamp."""
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if conversation:
|
|||
|
|
conversation.updated_at = datetime.now(timezone.utc)
|
|||
|
|
conversation.set_audit_fields(user_id=conversation.user_id, is_update=True)
|
|||
|
|
await self.session.commit()
|
|||
|
|
|
|||
|
|
async def get_user_conversations_count(
|
|||
|
|
self,
|
|||
|
|
search_query: Optional[str] = None,
|
|||
|
|
include_archived: bool = False
|
|||
|
|
) -> int:
|
|||
|
|
"""Get total count of user's conversations."""
|
|||
|
|
user_id = UserContext.get_current_user_id()
|
|||
|
|
query = select(func.count(Conversation.id)).where(
|
|||
|
|
Conversation.user_id == user_id
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
if not include_archived:
|
|||
|
|
query = query.where(Conversation.is_archived == False)
|
|||
|
|
|
|||
|
|
if search_query and search_query.strip():
|
|||
|
|
search_term = f"%{search_query.strip()}%"
|
|||
|
|
query = query.where(
|
|||
|
|
or_(
|
|||
|
|
Conversation.title.ilike(search_term),
|
|||
|
|
Conversation.system_prompt.ilike(search_term)
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
return (await self.session.scalar(query)) or 0
|
|||
|
|
|
|||
|
|
async def archive_conversation(self, conversation_id: int) -> bool:
|
|||
|
|
"""Archive a conversation."""
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if not conversation:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
conversation.is_archived = True
|
|||
|
|
conversation.set_audit_fields(user_id=conversation.user_id, is_update=True)
|
|||
|
|
await self.session.commit()
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
async def unarchive_conversation(self, conversation_id: int) -> bool:
|
|||
|
|
"""Unarchive a conversation."""
|
|||
|
|
conversation = await self.get_conversation(conversation_id)
|
|||
|
|
if not conversation:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
conversation.is_archived = False
|
|||
|
|
conversation.set_audit_fields(user_id=conversation.user_id, is_update=True)
|
|||
|
|
await self.session.commit()
|
|||
|
|
return True
|