Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,413 @@
|
||||
"""
|
||||
对话服务
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
from typing import AsyncGenerator, Dict, List, Optional
|
||||
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ai_platform.models import AIApp, Conversation, Message, LLMModel
|
||||
from utils.context import get_current_user_id_from_context
|
||||
from .llm_service import LLMService
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChatService:
|
||||
"""
|
||||
对话服务
|
||||
|
||||
管理对话和消息
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
self._db = db
|
||||
self.llm_service = LLMService(db)
|
||||
|
||||
async def create_conversation(
|
||||
self,
|
||||
app_id: str,
|
||||
title: str = '',
|
||||
user_id: Optional[str] = None,
|
||||
) -> Conversation:
|
||||
"""
|
||||
创建对话
|
||||
|
||||
Args:
|
||||
app_id: 应用 ID
|
||||
title: 对话标题
|
||||
|
||||
Returns:
|
||||
Conversation
|
||||
"""
|
||||
current_user_id = user_id or get_current_user_id_from_context()
|
||||
if not current_user_id:
|
||||
raise ValueError('未登录或登录已过期')
|
||||
|
||||
# 查询应用
|
||||
result = await self._db.execute(
|
||||
select(AIApp).where(AIApp.id == app_id, AIApp.is_deleted == False)
|
||||
)
|
||||
app = result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise ValueError(f'应用不存在: {app_id}')
|
||||
|
||||
conversation = Conversation(
|
||||
app_id=app_id,
|
||||
user_id=current_user_id,
|
||||
title=title or '新对话',
|
||||
)
|
||||
self._db.add(conversation)
|
||||
|
||||
# 更新应用统计
|
||||
app.conversation_count = (app.conversation_count or 0) + 1
|
||||
|
||||
await self._db.commit()
|
||||
await self._db.refresh(conversation)
|
||||
|
||||
return conversation
|
||||
|
||||
async def get_conversation(self, conversation_id: str) -> Optional[Conversation]:
|
||||
"""获取对话"""
|
||||
result = await self._db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list_conversations(
|
||||
self,
|
||||
app_id: str,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
) -> tuple:
|
||||
"""
|
||||
获取对话列表
|
||||
|
||||
Returns:
|
||||
(conversations, total)
|
||||
"""
|
||||
query = select(Conversation).where(
|
||||
Conversation.app_id == app_id,
|
||||
Conversation.is_deleted == False
|
||||
).order_by(Conversation.is_pinned.desc(), Conversation.sys_update_datetime.desc())
|
||||
|
||||
# 获取总数
|
||||
count_result = await self._db.execute(
|
||||
select(func.count()).select_from(query.subquery())
|
||||
)
|
||||
total = count_result.scalar() or 0
|
||||
|
||||
# 分页
|
||||
offset = (page - 1) * page_size
|
||||
query = query.offset(offset).limit(page_size)
|
||||
|
||||
result = await self._db.execute(query)
|
||||
conversations = result.scalars().all()
|
||||
|
||||
return list(conversations), total
|
||||
|
||||
async def delete_conversation(self, conversation_id: str) -> bool:
|
||||
"""删除对话"""
|
||||
result = await self._db.execute(
|
||||
select(Conversation).where(
|
||||
Conversation.id == conversation_id,
|
||||
Conversation.is_deleted == False
|
||||
)
|
||||
)
|
||||
conversation = result.scalar_one_or_none()
|
||||
if not conversation:
|
||||
return False
|
||||
|
||||
conversation.is_deleted = True
|
||||
await self._db.commit()
|
||||
return True
|
||||
|
||||
async def get_messages(
|
||||
self,
|
||||
conversation_id: str,
|
||||
limit: int = 50,
|
||||
) -> List[Message]:
|
||||
"""获取对话消息"""
|
||||
result = await self._db.execute(
|
||||
select(Message).where(
|
||||
Message.conversation_id == conversation_id,
|
||||
Message.is_deleted == False
|
||||
).order_by(Message.sys_create_datetime).limit(limit)
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def send_message(
|
||||
self,
|
||||
conversation_id: str,
|
||||
content: str,
|
||||
) -> tuple:
|
||||
"""
|
||||
发送消息并获取 AI 回复
|
||||
|
||||
Args:
|
||||
conversation_id: 对话 ID
|
||||
content: 消息内容
|
||||
|
||||
Returns:
|
||||
(user_message, assistant_message)
|
||||
"""
|
||||
conversation = await self.get_conversation(conversation_id)
|
||||
if not conversation:
|
||||
raise ValueError('对话不存在')
|
||||
|
||||
# 获取应用
|
||||
app_result = await self._db.execute(
|
||||
select(AIApp).where(AIApp.id == conversation.app_id)
|
||||
)
|
||||
app = app_result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise ValueError('应用不存在')
|
||||
|
||||
# 获取模型
|
||||
model = await self._get_effective_model(conversation, app)
|
||||
if not model:
|
||||
raise ValueError('未配置模型')
|
||||
|
||||
# 创建用户消息
|
||||
user_message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role='user',
|
||||
content=content,
|
||||
status='completed',
|
||||
)
|
||||
self._db.add(user_message)
|
||||
await self._db.flush()
|
||||
|
||||
# 构建消息列表
|
||||
messages = await self._build_messages(conversation, app, content)
|
||||
|
||||
# 创建助手消息（pending 状态）
|
||||
assistant_message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role='assistant',
|
||||
content='',
|
||||
status='pending',
|
||||
model_name=model.model_name,
|
||||
)
|
||||
self._db.add(assistant_message)
|
||||
await self._db.flush()
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
|
||||
# 调用 LLM
|
||||
response = await self.llm_service.chat_async(
|
||||
model_id=str(model.id),
|
||||
messages=messages,
|
||||
temperature=app.temperature or 0.7,
|
||||
max_tokens=app.max_tokens or 2048,
|
||||
)
|
||||
|
||||
latency = int((time.time() - start_time) * 1000)
|
||||
|
||||
# 更新助手消息
|
||||
assistant_message.content = response.content
|
||||
assistant_message.status = 'completed'
|
||||
assistant_message.prompt_tokens = response.prompt_tokens
|
||||
assistant_message.completion_tokens = response.completion_tokens
|
||||
assistant_message.total_tokens = response.total_tokens
|
||||
assistant_message.latency = latency
|
||||
|
||||
# 更新对话统计
|
||||
conversation.message_count = (conversation.message_count or 0) + 2
|
||||
conversation.total_tokens = (conversation.total_tokens or 0) + response.total_tokens
|
||||
|
||||
# 更新应用统计
|
||||
app.message_count = (app.message_count or 0) + 2
|
||||
|
||||
# 自动生成标题
|
||||
if conversation.message_count == 2:
|
||||
self._generate_title(conversation, content)
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
return user_message, assistant_message
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'发送消息失败: {e}')
|
||||
assistant_message.status = 'failed'
|
||||
assistant_message.error_message = str(e)
|
||||
await self._db.commit()
|
||||
raise
|
||||
|
||||
async def send_message_stream(
|
||||
self,
|
||||
conversation_id: str,
|
||||
content: str,
|
||||
) -> AsyncGenerator[Dict, None]:
|
||||
"""
|
||||
流式发送消息
|
||||
|
||||
Yields:
|
||||
{"type": "content", "content": "..."} 或
|
||||
{"type": "done", "message": {...}}
|
||||
"""
|
||||
conversation = await self.get_conversation(conversation_id)
|
||||
if not conversation:
|
||||
raise ValueError('对话不存在')
|
||||
|
||||
# 获取应用
|
||||
app_result = await self._db.execute(
|
||||
select(AIApp).where(AIApp.id == conversation.app_id)
|
||||
)
|
||||
app = app_result.scalar_one_or_none()
|
||||
if not app:
|
||||
raise ValueError('应用不存在')
|
||||
|
||||
# 获取模型
|
||||
model = await self._get_effective_model(conversation, app)
|
||||
if not model:
|
||||
raise ValueError('未配置模型')
|
||||
|
||||
# 创建用户消息
|
||||
user_message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role='user',
|
||||
content=content,
|
||||
status='completed',
|
||||
)
|
||||
self._db.add(user_message)
|
||||
await self._db.flush()
|
||||
|
||||
# 构建消息列表
|
||||
messages = await self._build_messages(conversation, app, content)
|
||||
|
||||
# 创建助手消息
|
||||
assistant_message = Message(
|
||||
conversation_id=conversation_id,
|
||||
role='assistant',
|
||||
content='',
|
||||
status='pending',
|
||||
model_name=model.model_name,
|
||||
)
|
||||
self._db.add(assistant_message)
|
||||
await self._db.flush()
|
||||
|
||||
try:
|
||||
start_time = time.time()
|
||||
full_content = ''
|
||||
total_tokens = 0
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
|
||||
# 使用异步流式方法
|
||||
async for chunk in self.llm_service.chat_stream(
|
||||
model_id=str(model.id),
|
||||
messages=messages,
|
||||
temperature=app.temperature or 0.7,
|
||||
max_tokens=app.max_tokens or 2048,
|
||||
):
|
||||
if chunk.content:
|
||||
full_content += chunk.content
|
||||
yield {'type': 'content', 'content': chunk.content}
|
||||
|
||||
if chunk.is_finished:
|
||||
prompt_tokens = chunk.prompt_tokens
|
||||
completion_tokens = chunk.completion_tokens
|
||||
total_tokens = chunk.total_tokens
|
||||
|
||||
latency = int((time.time() - start_time) * 1000)
|
||||
|
||||
# 更新助手消息
|
||||
assistant_message.content = full_content
|
||||
assistant_message.status = 'completed'
|
||||
assistant_message.prompt_tokens = prompt_tokens
|
||||
assistant_message.completion_tokens = completion_tokens
|
||||
assistant_message.total_tokens = total_tokens
|
||||
assistant_message.latency = latency
|
||||
|
||||
# 更新对话统计
|
||||
conversation.message_count = (conversation.message_count or 0) + 2
|
||||
conversation.total_tokens = (conversation.total_tokens or 0) + total_tokens
|
||||
|
||||
await self._db.commit()
|
||||
|
||||
yield {
|
||||
'type': 'done',
|
||||
'message': {
|
||||
'id': str(assistant_message.id),
|
||||
'content': full_content,
|
||||
'tokens': total_tokens,
|
||||
'latency': latency,
|
||||
},
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'流式发送消息失败: {e}')
|
||||
assistant_message.status = 'failed'
|
||||
assistant_message.error_message = str(e)
|
||||
await self._db.commit()
|
||||
yield {'type': 'error', 'error': str(e)}
|
||||
|
||||
async def _get_effective_model(self, conversation: Conversation, app: AIApp) -> Optional[LLMModel]:
|
||||
"""获取有效的模型"""
|
||||
model_id = conversation.model_override_id or app.model_id
|
||||
if not model_id:
|
||||
return None
|
||||
|
||||
result = await self._db.execute(
|
||||
select(LLMModel).where(
|
||||
LLMModel.id == model_id,
|
||||
LLMModel.is_active == True,
|
||||
LLMModel.is_deleted == False
|
||||
)
|
||||
)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def _build_messages(
|
||||
self,
|
||||
conversation: Conversation,
|
||||
app: AIApp,
|
||||
user_content: str,
|
||||
) -> List[Dict[str, str]]:
|
||||
"""构建消息列表"""
|
||||
messages = []
|
||||
|
||||
# 系统提示词
|
||||
if app.system_prompt:
|
||||
messages.append({
|
||||
'role': 'system',
|
||||
'content': app.system_prompt,
|
||||
})
|
||||
|
||||
# 历史消息
|
||||
result = await self._db.execute(
|
||||
select(Message).where(
|
||||
Message.conversation_id == conversation.id,
|
||||
Message.is_deleted == False,
|
||||
Message.status == 'completed'
|
||||
).order_by(Message.sys_create_datetime.desc()).limit(20)
|
||||
)
|
||||
history = result.scalars().all()
|
||||
|
||||
for msg in reversed(list(history)):
|
||||
messages.append({
|
||||
'role': msg.role,
|
||||
'content': msg.content,
|
||||
})
|
||||
|
||||
# 当前用户消息
|
||||
messages.append({
|
||||
'role': 'user',
|
||||
'content': user_content,
|
||||
})
|
||||
|
||||
return messages
|
||||
|
||||
def _generate_title(self, conversation: Conversation, first_message: str):
|
||||
"""自动生成对话标题"""
|
||||
# 简单截取前 20 个字符作为标题
|
||||
title = first_message[:20]
|
||||
if len(first_message) > 20:
|
||||
title += '...'
|
||||
conversation.title = title
|
||||
Reference in New Issue
Block a user