1631 lines
60 KiB
Python
1631 lines
60 KiB
Python
"""
|
||
试运行服务
|
||
|
||
提供 Agent 试运行功能，允许用户在不发布应用的情况下测试配置。
|
||
"""
|
||
import time
|
||
import uuid
|
||
import json
|
||
import asyncio
|
||
import datetime
|
||
from typing import Dict, Any, Optional, List, AsyncGenerator
|
||
from langchain.tools import tool
|
||
from pydantic import BaseModel, Field
|
||
from sqlalchemy.orm import Session
|
||
from sqlalchemy import select
|
||
|
||
from app.services.memory_konwledges_server import write_rag
|
||
from app.tasks import write_message_task
|
||
from app.models import AgentConfig, ModelConfig, ModelApiKey
|
||
from app.core.exceptions import BusinessException, ResourceNotFoundException
|
||
from app.core.error_codes import BizCode
|
||
from app.core.logging_config import get_business_logger
|
||
from app.schemas.prompt_schema import render_prompt_message, PromptMessageRole
|
||
from app.services.memory_agent_service import MemoryAgentService
|
||
from app.services.model_parameter_merger import ModelParameterMerger
|
||
from app.core.rag.nlp.search import knowledge_retrieval
|
||
from app.services.langchain_tool_server import Search
|
||
from app.services.task_service import get_task_memory_write_result
|
||
|
||
logger = get_business_logger()
|
||
class KnowledgeRetrievalInput(BaseModel):
|
||
"""知识库检索工具输入参数"""
|
||
query: str = Field(description="需要检索的问题或关键词")
|
||
|
||
|
||
class WebSearchInput(BaseModel):
|
||
"""网络搜索工具输入参数"""
|
||
query: str = Field(description="需要搜索的问题或关键词")
|
||
|
||
|
||
class LongTermMemoryInput(BaseModel):
|
||
"""长期记忆工具输入参数"""
|
||
question: str = Field(description="需要查询的问题")
|
||
|
||
def create_long_term_memory_tool(memory_config: Dict[str, Any], end_user_id: str, storage_type: Optional[str] = None,user_rag_memory_id: Optional[str] = None):
|
||
"""创建长期记忆工具
|
||
|
||
Args:
|
||
memory_config: 记忆配置
|
||
end_user_id: 用户ID
|
||
storage_type: 存储类型（可选）
|
||
|
||
Returns:
|
||
长期记忆工具
|
||
"""
|
||
# search_switch = memory_config.get("search_switch", "2")
|
||
config_id= memory_config.get("memory_content",'17')
|
||
|
||
logger.info(f"创建长期记忆工具，配置: end_user_id={end_user_id}, config_id={config_id}, storage_type={storage_type}")
|
||
|
||
@tool(args_schema=LongTermMemoryInput)
|
||
def long_term_memory(question: str) -> str:
|
||
"""从长期记忆中检索历史对话信息。当需要回忆之前的对话内容、用户偏好或历史信息时使用此工具。
|
||
|
||
Args:
|
||
question: 需要查询的问题
|
||
end_user_id: 用户唯一标识符
|
||
search_switch: 搜索开关（on/off）
|
||
|
||
Returns:
|
||
检索到的历史记忆内容
|
||
"""
|
||
logger.info(f" 长期记忆工具被调用！question={question}, user={end_user_id}")
|
||
|
||
try:
|
||
memory_content = asyncio.run(
|
||
MemoryAgentService().read_memory(
|
||
group_id=end_user_id,
|
||
message=question,
|
||
history=[],
|
||
search_switch="2",
|
||
config_id=config_id,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id
|
||
)
|
||
)
|
||
logger.info(f'用户ID：Agent:{end_user_id}')
|
||
logger.debug(f"调用长期记忆 API", extra={"question": question, "end_user_id": end_user_id})
|
||
|
||
logger.info(
|
||
f"长期记忆检索成功",
|
||
extra={
|
||
"end_user_id": end_user_id,
|
||
"content_length": len(str(memory_content))
|
||
}
|
||
)
|
||
|
||
return f"检索到以下历史记忆：\n\n{memory_content}"
|
||
except Exception as e:
|
||
logger.error(f"长期记忆检索失败", extra={"error": str(e), "error_type": type(e).__name__})
|
||
return f"记忆检索失败: {str(e)}"
|
||
|
||
return long_term_memory
|
||
|
||
|
||
def create_web_search_tool(web_search_config: Dict[str, Any]):
|
||
"""创建网络搜索工具
|
||
|
||
Args:
|
||
web_search_config: 网络搜索配置
|
||
|
||
Returns:
|
||
网络搜索工具
|
||
"""
|
||
logger.info("创建网络搜索工具")
|
||
|
||
@tool(args_schema=WebSearchInput)
|
||
def web_search_tool(query: str) -> str:
|
||
"""从互联网搜索最新信息。当用户的问题需要实时信息、最新新闻或网络资料时，使用此工具进行搜索。
|
||
|
||
Args:
|
||
query: 需要搜索的问题或关键词
|
||
|
||
Returns:
|
||
搜索到的相关网络信息
|
||
"""
|
||
try:
|
||
logger.info(f"执行网络搜索: {query}")
|
||
|
||
# 调用搜索服务
|
||
search_result = Search(query)
|
||
logger.info(
|
||
"网络搜索成功",
|
||
extra={
|
||
"query": query,
|
||
"result_length": len(search_result)
|
||
}
|
||
)
|
||
|
||
return f"搜索到以下网络信息：\n\n{search_result}"
|
||
|
||
except Exception as e:
|
||
logger.error(f"网络搜索失败", extra={"error": str(e), "error_type": type(e).__name__})
|
||
return f"搜索失败: {str(e)}"
|
||
|
||
return web_search_tool
|
||
|
||
|
||
def create_knowledge_retrieval_tool(kb_config,kb_ids,user_id):
|
||
"""从知识库中检索相关信息。当用户的问题需要参考知识库、文档或历史记录时，使用此工具进行检索。
|
||
|
||
Args:
|
||
query: 需要检索的问题或关键词
|
||
|
||
Returns:
|
||
检索到的相关知识内容
|
||
"""
|
||
logger.info(f"创建知识库检索工具，用户：{user_id}")
|
||
@tool(args_schema=KnowledgeRetrievalInput)
|
||
def knowledge_retrieval_tool(query: str) -> str:
|
||
"""从知识库中检索相关信息。当用户的问题需要参考知识库、文档或历史记录时，使用此工具进行检索。
|
||
|
||
Args:
|
||
query: 需要检索的问题或关键词
|
||
|
||
Returns:
|
||
检索到的相关知识内容
|
||
"""
|
||
|
||
|
||
try:
|
||
|
||
retrieve_chunks_result = knowledge_retrieval(query, kb_config)
|
||
if retrieve_chunks_result:
|
||
retrieval_knowledge = [i.page_content for i in retrieve_chunks_result]
|
||
context = '\n\n'.join(retrieval_knowledge)
|
||
logger.info(
|
||
f"知识库检索成功",
|
||
extra={
|
||
"kb_ids": kb_ids,
|
||
"result_count": len(retrieval_knowledge),
|
||
"total_length": len(context)
|
||
}
|
||
)
|
||
|
||
return f"检索到以下相关信息：\n\n{context}"
|
||
else:
|
||
logger.warning("知识库检索未找到结果")
|
||
return "未找到相关信息"
|
||
except Exception as e:
|
||
logger.error(f"知识库检索失败", extra={"error": str(e), "error_type": type(e).__name__})
|
||
return f"检索失败: {str(e)}"
|
||
|
||
return knowledge_retrieval_tool
|
||
|
||
class DraftRunService:
|
||
"""试运行服务类"""
|
||
|
||
def __init__(self, db: Session):
|
||
"""初始化试运行服务
|
||
|
||
Args:
|
||
db: 数据库会话
|
||
"""
|
||
self.db = db
|
||
|
||
async def run(
|
||
self,
|
||
*,
|
||
agent_config: AgentConfig,
|
||
model_config: ModelConfig,
|
||
message: str,
|
||
workspace_id: uuid.UUID,
|
||
conversation_id: Optional[str] = None,
|
||
user_id: Optional[str] = None,
|
||
variables: Optional[Dict[str, Any]] = None,
|
||
storage_type: Optional[str] = None,
|
||
user_rag_memory_id: Optional[str] = None,
|
||
web_search: bool = True,
|
||
memory: bool = True
|
||
) -> Dict[str, Any]:
|
||
"""执行试运行（使用 LangChain Agent）
|
||
|
||
Args:
|
||
agent_config: Agent 配置
|
||
model_config: 模型配置
|
||
message: 用户消息
|
||
workspace_id: 工作空间ID（必须，用于会话隔离）
|
||
conversation_id: 会话ID（用于多轮对话）
|
||
user_id: 用户ID
|
||
variables: 自定义变量参数值
|
||
|
||
Returns:
|
||
Dict: 包含 AI 回复和元数据的字典
|
||
"""
|
||
|
||
print('===========',storage_type)
|
||
|
||
print(user_id)
|
||
if variables == None: variables = {}
|
||
from app.core.agent.langchain_agent import LangChainAgent
|
||
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# 1. 获取 API Key 配置
|
||
api_key_config = await self._get_api_key(model_config.id)
|
||
logger.debug(
|
||
f"API Key 配置获取成功",
|
||
extra={
|
||
"model_name": api_key_config["model_name"],
|
||
"has_api_key": bool(api_key_config["api_key"]),
|
||
"has_api_base": bool(api_key_config.get("api_base"))
|
||
}
|
||
)
|
||
|
||
# 2. 合并模型参数
|
||
effective_params = ModelParameterMerger.get_effective_parameters(
|
||
model_config=model_config,
|
||
agent_config=agent_config
|
||
)
|
||
|
||
|
||
items_params=variables
|
||
system_prompt = render_prompt_message(
|
||
agent_config.system_prompt, # 修正拼写错误
|
||
PromptMessageRole.USER,
|
||
items_params
|
||
)
|
||
|
||
# 3. 处理系统提示词（支持变量替换）
|
||
system_prompt = system_prompt.get_text_content() or "你是一个专业的AI助手"
|
||
print('系统提示词：',system_prompt)
|
||
|
||
# 4. 准备工具列表
|
||
tools = []
|
||
|
||
# 添加网络搜索工具
|
||
if web_search:
|
||
if agent_config.tools:
|
||
web_search_config = agent_config.tools.get("web_search", {})
|
||
web_search_enable = web_search_config.get("enabled", False)
|
||
|
||
if web_search_enable:
|
||
logger.info("网络搜索已启用")
|
||
# 创建网络搜索工具
|
||
search_tool = create_web_search_tool(web_search_config)
|
||
tools.append(search_tool)
|
||
|
||
logger.debug(
|
||
"已添加网络搜索工具",
|
||
extra={
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 添加知识库检索工具
|
||
if agent_config.knowledge_retrieval:
|
||
kb_config = agent_config.knowledge_retrieval
|
||
knowledge_bases = kb_config.get("knowledge_bases", [])
|
||
kb_ids = bool(knowledge_bases and knowledge_bases[0].get("kb_id"))
|
||
if kb_ids:
|
||
# 创建知识库检索工具
|
||
kb_tool = create_knowledge_retrieval_tool(kb_config,kb_ids,user_id)
|
||
tools.append(kb_tool)
|
||
|
||
logger.debug(
|
||
f"已添加知识库检索工具",
|
||
extra={
|
||
"kb_ids": kb_ids,
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 添加长期记忆工具
|
||
if memory:
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
|
||
memory_config = agent_config.memory
|
||
if user_id:
|
||
# 创建长期记忆工具
|
||
memory_tool = create_long_term_memory_tool(memory_config, user_id,storage_type,user_rag_memory_id)
|
||
tools.append(memory_tool)
|
||
|
||
logger.debug(
|
||
f"已添加长期记忆工具",
|
||
extra={
|
||
"user_id": user_id,
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 4. 创建 LangChain Agent
|
||
agent = LangChainAgent(
|
||
model_name=api_key_config["model_name"],
|
||
api_key=api_key_config["api_key"],
|
||
provider=api_key_config.get("provider", "openai"),
|
||
api_base=api_key_config.get("api_base"),
|
||
temperature=effective_params.get("temperature", 0.7),
|
||
max_tokens=effective_params.get("max_tokens", 2000),
|
||
system_prompt=system_prompt,
|
||
tools=tools,
|
||
)
|
||
|
||
# 5. 处理会话ID（创建或验证）
|
||
conversation_id = await self._ensure_conversation(
|
||
conversation_id=conversation_id,
|
||
app_id=agent_config.app_id,
|
||
workspace_id=workspace_id,
|
||
user_id=user_id
|
||
)
|
||
|
||
# 6. 加载历史消息
|
||
history = []
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
history = await self._load_conversation_history(
|
||
conversation_id=conversation_id,
|
||
max_history=agent_config.memory.get("max_history", 10)
|
||
)
|
||
|
||
# 6. 知识库检索
|
||
context = None
|
||
|
||
logger.debug(
|
||
f"准备调用 LangChain Agent",
|
||
extra={
|
||
"model": api_key_config["model_name"],
|
||
"has_history": bool(history),
|
||
"has_context": bool(context)
|
||
}
|
||
)
|
||
|
||
memory_config_= agent_config.memory
|
||
config_id = memory_config_.get("memory_content")
|
||
|
||
# 7. 调用 Agent
|
||
result = await agent.chat(
|
||
message=message,
|
||
history=history,
|
||
context=context,
|
||
end_user_id=user_id,
|
||
config_id=config_id,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id
|
||
)
|
||
|
||
elapsed_time = time.time() - start_time
|
||
|
||
# 8. 保存会话消息
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
await self._save_conversation_message(
|
||
conversation_id=conversation_id,
|
||
user_message=message,
|
||
assistant_message=result["content"],
|
||
app_id=agent_config.app_id,
|
||
user_id=user_id
|
||
)
|
||
|
||
response = {
|
||
"message": result["content"],
|
||
"conversation_id": conversation_id,
|
||
"usage": result.get("usage", {
|
||
"prompt_tokens": 0,
|
||
"completion_tokens": 0,
|
||
"total_tokens": 0
|
||
}),
|
||
"elapsed_time": elapsed_time
|
||
}
|
||
|
||
logger.info(
|
||
f"试运行完成",
|
||
extra={
|
||
"model": model_config.name,
|
||
"elapsed_time": elapsed_time,
|
||
"message_length": len(result["content"]),
|
||
"total_tokens": result.get("usage", {}).get("total_tokens", 0)
|
||
}
|
||
)
|
||
|
||
return response
|
||
|
||
except Exception as e:
|
||
logger.error(f"LangChain Agent 调用失败", extra={"error": str(e), "error_type": type(e).__name__})
|
||
raise BusinessException(f"Agent 调用失败: {str(e)}", BizCode.INTERNAL_ERROR, cause=e)
|
||
|
||
async def run_stream(
|
||
self,
|
||
*,
|
||
agent_config: AgentConfig,
|
||
model_config: ModelConfig,
|
||
message: str,
|
||
workspace_id: uuid.UUID,
|
||
conversation_id: Optional[str] = None,
|
||
user_id: Optional[str] = None,
|
||
variables: Optional[Dict[str, Any]] = None,
|
||
storage_type: Optional[str] = None,
|
||
user_rag_memory_id: Optional[str] = None,
|
||
web_search: bool = True, # 布尔类型默认值
|
||
memory: bool = True # 布尔类型默认值
|
||
|
||
) -> AsyncGenerator[str, None]:
|
||
"""执行试运行（流式返回，使用 LangChain Agent）
|
||
|
||
Args:
|
||
agent_config: Agent 配置
|
||
model_config: 模型配置
|
||
message: 用户消息
|
||
workspace_id: 工作空间ID（必须，用于会话隔离）
|
||
conversation_id: 会话ID（用于多轮对话）
|
||
user_id: 用户ID
|
||
variables: 自定义变量参数值
|
||
|
||
Yields:
|
||
str: SSE 格式的事件数据
|
||
"""
|
||
if variables==None:variables={}
|
||
|
||
from app.core.agent.langchain_agent import LangChainAgent
|
||
|
||
start_time = time.time()
|
||
|
||
try:
|
||
# 1. 获取 API Key 配置
|
||
api_key_config = await self._get_api_key(model_config.id)
|
||
|
||
# 2. 合并模型参数
|
||
effective_params = ModelParameterMerger.get_effective_parameters(
|
||
model_config=model_config,
|
||
agent_config=agent_config
|
||
)
|
||
|
||
items_params=variables
|
||
|
||
system_prompt = render_prompt_message(
|
||
agent_config.system_prompt, # 修正拼写错误
|
||
PromptMessageRole.USER,
|
||
items_params
|
||
)
|
||
|
||
# 3. 处理系统提示词（支持变量替换）
|
||
system_prompt = system_prompt.get_text_content() or "你是一个专业的AI助手"
|
||
|
||
# 4. 准备工具列表
|
||
tools = []
|
||
|
||
# 添加网络搜索工具
|
||
if web_search:
|
||
if agent_config.tools:
|
||
web_search = agent_config.tools.get("web_search", {})
|
||
web_search_enable = web_search.get("enable", False)
|
||
|
||
if web_search_enable:
|
||
logger.info("网络搜索已启用（流式）")
|
||
# 创建网络搜索工具
|
||
search_tool = create_web_search_tool(web_search)
|
||
tools.append(search_tool)
|
||
|
||
logger.debug(
|
||
"已添加网络搜索工具（流式）",
|
||
extra={
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 添加知识库检索工具
|
||
if agent_config.knowledge_retrieval:
|
||
kb_config = agent_config.knowledge_retrieval
|
||
knowledge_bases = kb_config.get("knowledge_bases", [])
|
||
kb_ids = bool(knowledge_bases and knowledge_bases[0].get("kb_id"))
|
||
if kb_ids:
|
||
# 创建知识库检索工具
|
||
kb_tool = create_knowledge_retrieval_tool(kb_config,kb_ids,user_id)
|
||
tools.append(kb_tool)
|
||
|
||
logger.debug(
|
||
f"已添加知识库检索工具",
|
||
extra={
|
||
"kb_ids": kb_ids,
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 添加长期记忆工具
|
||
if memory:
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
memory_config = agent_config.memory
|
||
if user_id:
|
||
# 创建长期记忆工具
|
||
memory_tool = create_long_term_memory_tool(memory_config, user_id,storage_type,user_rag_memory_id)
|
||
tools.append(memory_tool)
|
||
|
||
logger.debug(
|
||
f"已添加长期记忆工具",
|
||
extra={
|
||
"user_id": user_id,
|
||
"tool_count": len(tools)
|
||
}
|
||
)
|
||
|
||
# 4. 创建 LangChain Agent
|
||
agent = LangChainAgent(
|
||
model_name=api_key_config["model_name"],
|
||
api_key=api_key_config["api_key"],
|
||
provider=api_key_config.get("provider", "openai"),
|
||
api_base=api_key_config.get("api_base"),
|
||
temperature=effective_params.get("temperature", 0.7),
|
||
max_tokens=effective_params.get("max_tokens", 2000),
|
||
system_prompt=system_prompt,
|
||
tools=tools,
|
||
streaming=True
|
||
)
|
||
|
||
# 5. 处理会话ID（创建或验证）
|
||
conversation_id = await self._ensure_conversation(
|
||
conversation_id=conversation_id,
|
||
app_id=agent_config.app_id,
|
||
workspace_id=workspace_id,
|
||
user_id=user_id
|
||
)
|
||
|
||
# 6. 加载历史消息
|
||
history = []
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
history = await self._load_conversation_history(
|
||
conversation_id=conversation_id,
|
||
max_history=agent_config.memory.get("max_history", 10)
|
||
)
|
||
|
||
# 7. 知识库检索
|
||
context = None
|
||
|
||
# 8. 发送开始事件
|
||
yield self._format_sse_event("start", {
|
||
"conversation_id": conversation_id,
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
memory_config_ = agent_config.memory
|
||
config_id = memory_config_.get("memory_content")
|
||
|
||
# 9. 流式调用 Agent
|
||
full_content = ""
|
||
async for chunk in agent.chat_stream(
|
||
message=message,
|
||
history=history,
|
||
context=context,
|
||
end_user_id=user_id,
|
||
config_id=config_id,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id
|
||
):
|
||
full_content += chunk
|
||
# 发送消息块事件
|
||
yield self._format_sse_event("message", {
|
||
"content": chunk
|
||
})
|
||
|
||
if storage_type == "rag":
|
||
await write_rag(user_id, full_content, user_rag_memory_id)
|
||
else:
|
||
write_id = write_message_task.delay(user_id, full_content, config_id, storage_type, user_rag_memory_id)
|
||
write_status = get_task_memory_write_result(str(write_id))
|
||
logger.info(f'Agent:{user_id};{full_content}--{write_status}')
|
||
|
||
elapsed_time = time.time() - start_time
|
||
|
||
# 10. 保存会话消息
|
||
if agent_config.memory and agent_config.memory.get("enabled"):
|
||
await self._save_conversation_message(
|
||
conversation_id=conversation_id,
|
||
user_message=message,
|
||
assistant_message=full_content,
|
||
app_id=agent_config.app_id,
|
||
user_id=user_id
|
||
)
|
||
|
||
# 11. 发送结束事件
|
||
yield self._format_sse_event("end", {
|
||
"conversation_id": conversation_id,
|
||
"elapsed_time": elapsed_time,
|
||
"message_length": len(full_content)
|
||
})
|
||
|
||
logger.info(
|
||
f"流式试运行完成",
|
||
extra={
|
||
"model": model_config.name,
|
||
"elapsed_time": elapsed_time,
|
||
"message_length": len(full_content)
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"流式 Agent 调用失败", extra={"error": str(e)})
|
||
# 发送错误事件
|
||
yield self._format_sse_event("error", {
|
||
"error": str(e),
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
def _format_sse_event(self, event: str, data: Dict[str, Any]) -> str:
|
||
"""格式化 SSE 事件
|
||
|
||
Args:
|
||
event: 事件类型
|
||
data: 事件数据
|
||
|
||
Returns:
|
||
str: SSE 格式的字符串
|
||
"""
|
||
return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||
|
||
async def _get_api_key(self, model_config_id: uuid.UUID) -> Dict[str, str]:
|
||
"""获取模型的 API Key
|
||
|
||
Args:
|
||
model_config_id: 模型配置ID
|
||
|
||
Returns:
|
||
Dict: 包含 model_name, api_key, api_base 的字典
|
||
|
||
Raises:
|
||
BusinessException: 当没有可用的 API Key 时
|
||
"""
|
||
stmt = (
|
||
select(ModelApiKey)
|
||
.where(
|
||
ModelApiKey.model_config_id == model_config_id,
|
||
ModelApiKey.is_active == True
|
||
)
|
||
.order_by(ModelApiKey.priority.desc())
|
||
.limit(1)
|
||
)
|
||
|
||
api_key = self.db.scalars(stmt).first()
|
||
|
||
if not api_key:
|
||
raise BusinessException("没有可用的 API Key", BizCode.AGENT_CONFIG_MISSING)
|
||
|
||
return {
|
||
"model_name": api_key.model_name,
|
||
"provider": api_key.provider,
|
||
"api_key": api_key.api_key,
|
||
"api_base": api_key.api_base
|
||
}
|
||
|
||
async def _ensure_conversation(
|
||
self,
|
||
conversation_id: Optional[str],
|
||
app_id: uuid.UUID,
|
||
workspace_id: uuid.UUID,
|
||
user_id: Optional[str]
|
||
) -> str:
|
||
"""确保会话存在（创建或验证）
|
||
|
||
Args:
|
||
conversation_id: 会话ID（可选）
|
||
app_id: 应用ID
|
||
workspace_id: 工作空间ID（必须）
|
||
user_id: 用户ID
|
||
|
||
Returns:
|
||
str: 会话ID
|
||
|
||
Raises:
|
||
BusinessException: 当指定的会话不存在时
|
||
"""
|
||
from app.services.conversation_service import ConversationService
|
||
from app.schemas.conversation_schema import ConversationCreate
|
||
from app.models import Conversation as ConversationModel
|
||
|
||
conversation_service = ConversationService(self.db)
|
||
|
||
# 如果没有提供会话ID，创建新会话
|
||
if not conversation_id:
|
||
logger.info(
|
||
"创建新的草稿会话",
|
||
extra={"workspace_id": str(workspace_id)}
|
||
)
|
||
|
||
# 获取配置快照
|
||
config_snapshot = await self._get_config_snapshot(app_id)
|
||
|
||
# 创建新会话
|
||
new_conv_id = str(uuid.uuid4())
|
||
new_conversation = ConversationModel(
|
||
id=uuid.UUID(new_conv_id),
|
||
app_id=app_id,
|
||
workspace_id=workspace_id,
|
||
user_id=user_id,
|
||
is_draft=True,
|
||
title="草稿会话",
|
||
config_snapshot=config_snapshot
|
||
)
|
||
self.db.add(new_conversation)
|
||
self.db.commit()
|
||
self.db.refresh(new_conversation)
|
||
|
||
logger.info(
|
||
f"创建草稿会话成功",
|
||
extra={
|
||
"conversation_id": new_conv_id,
|
||
"workspace_id": str(workspace_id)
|
||
}
|
||
)
|
||
|
||
return new_conv_id
|
||
|
||
# 如果提供了会话ID，验证其存在性和工作空间归属
|
||
try:
|
||
conv_uuid = uuid.UUID(conversation_id)
|
||
conversation = conversation_service.get_conversation(conv_uuid)
|
||
|
||
# 验证会话属于当前工作空间
|
||
if conversation.workspace_id != workspace_id:
|
||
logger.warning(
|
||
f"会话不属于当前工作空间",
|
||
extra={
|
||
"conversation_id": conversation_id,
|
||
"conversation_workspace_id": str(conversation.workspace_id),
|
||
"current_workspace_id": str(workspace_id)
|
||
}
|
||
)
|
||
raise BusinessException(
|
||
f"会话不属于当前工作空间",
|
||
BizCode.PERMISSION_DENIED
|
||
)
|
||
|
||
logger.debug(
|
||
f"使用现有会话",
|
||
extra={
|
||
"conversation_id": conversation_id,
|
||
"workspace_id": str(workspace_id)
|
||
}
|
||
)
|
||
return conversation_id
|
||
except BusinessException:
|
||
raise
|
||
except Exception as e:
|
||
logger.error(
|
||
f"会话不存在或无效",
|
||
extra={"conversation_id": conversation_id, "error": str(e)}
|
||
)
|
||
raise BusinessException(
|
||
f"会话不存在: {conversation_id}",
|
||
BizCode.NOT_FOUND,
|
||
cause=e
|
||
)
|
||
|
||
async def _load_conversation_history(
|
||
self,
|
||
conversation_id: str,
|
||
max_history: int = 10
|
||
) -> List[Dict[str, str]]:
|
||
"""加载会话历史消息
|
||
|
||
Args:
|
||
conversation_id: 会话ID
|
||
max_history: 最大历史消息数量
|
||
|
||
Returns:
|
||
List[Dict]: 历史消息列表
|
||
"""
|
||
try:
|
||
from app.services.conversation_service import ConversationService
|
||
|
||
conversation_service = ConversationService(self.db)
|
||
history = conversation_service.get_conversation_history(
|
||
conversation_id=uuid.UUID(conversation_id),
|
||
max_history=max_history
|
||
)
|
||
|
||
logger.debug(
|
||
f"加载会话历史",
|
||
extra={
|
||
"conversation_id": conversation_id,
|
||
"max_history": max_history,
|
||
"loaded_count": len(history)
|
||
}
|
||
)
|
||
|
||
return history
|
||
|
||
except Exception as e:
|
||
# 新会话没有历史记录是正常的
|
||
logger.debug(f"加载会话历史失败（可能是新会话）", extra={"error": str(e)})
|
||
return []
|
||
|
||
async def _save_conversation_message(
|
||
self,
|
||
conversation_id: str,
|
||
user_message: str,
|
||
assistant_message: str,
|
||
app_id: Optional[uuid.UUID] = None,
|
||
user_id: Optional[str] = None
|
||
) -> None:
|
||
"""保存会话消息（会话已通过 _ensure_conversation 确保存在）
|
||
|
||
Args:
|
||
conversation_id: 会话ID
|
||
user_message: 用户消息
|
||
assistant_message: AI 回复消息
|
||
app_id: 应用ID（未使用，保留用于兼容性）
|
||
user_id: 用户ID（未使用，保留用于兼容性）
|
||
"""
|
||
try:
|
||
from app.services.conversation_service import ConversationService
|
||
|
||
conversation_service = ConversationService(self.db)
|
||
conv_uuid = uuid.UUID(conversation_id)
|
||
|
||
# 保存消息（会话已经存在）
|
||
# 保存用户消息
|
||
conversation_service.add_message(
|
||
conversation_id=conv_uuid,
|
||
role="user",
|
||
content=user_message
|
||
)
|
||
# 保存助手消息
|
||
conversation_service.add_message(
|
||
conversation_id=conv_uuid,
|
||
role="assistant",
|
||
content=assistant_message
|
||
)
|
||
|
||
logger.debug(
|
||
f"保存会话消息",
|
||
extra={
|
||
"conversation_id": conversation_id,
|
||
"user_message_length": len(user_message),
|
||
"assistant_message_length": len(assistant_message)
|
||
}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.warning(f"保存会话消息失败", extra={"error": str(e)})
|
||
|
||
async def _get_config_snapshot(self, app_id: uuid.UUID) -> Dict[str, Any]:
|
||
"""获取当前配置快照
|
||
|
||
Args:
|
||
app_id: 应用ID
|
||
|
||
Returns:
|
||
Dict: 配置快照
|
||
"""
|
||
try:
|
||
from app.models import AgentConfig, ModelConfig
|
||
|
||
# 获取 Agent 配置
|
||
stmt = select(AgentConfig).where(AgentConfig.app_id == app_id)
|
||
agent_cfg = self.db.scalars(stmt).first()
|
||
|
||
if not agent_cfg:
|
||
return {}
|
||
|
||
# 获取模型配置
|
||
model_config = None
|
||
if agent_cfg.default_model_config_id:
|
||
model_config = self.db.get(ModelConfig, agent_cfg.default_model_config_id)
|
||
|
||
# 构建快照（确保所有值都可序列化）
|
||
def safe_serialize(value):
|
||
"""安全序列化值"""
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, (str, int, float, bool)):
|
||
return value
|
||
if isinstance(value, (dict, list)):
|
||
return value
|
||
# 对于 Pydantic 模型或其他对象，尝试转换为字典
|
||
if hasattr(value, 'dict'):
|
||
return value.dict()
|
||
if hasattr(value, '__dict__'):
|
||
return value.__dict__
|
||
return str(value)
|
||
|
||
snapshot = {
|
||
"agent_config": {
|
||
"system_prompt": agent_cfg.system_prompt,
|
||
"model_parameters": safe_serialize(agent_cfg.model_parameters),
|
||
"knowledge_retrieval": safe_serialize(agent_cfg.knowledge_retrieval),
|
||
"memory": safe_serialize(agent_cfg.memory),
|
||
"variables": safe_serialize(agent_cfg.variables),
|
||
"tools": safe_serialize(agent_cfg.tools)
|
||
},
|
||
"model_config": {
|
||
"model_name": model_config.name if model_config else None,
|
||
"provider": model_config.provider if model_config else None,
|
||
"type": model_config.type if model_config else None
|
||
} if model_config else None,
|
||
"snapshot_time": datetime.datetime.now().isoformat()
|
||
}
|
||
|
||
return snapshot
|
||
|
||
except Exception as e:
|
||
# 对于多 Agent 应用，没有直接的 AgentConfig 是正常的
|
||
logger.debug(f"获取配置快照失败（可能是多 Agent 应用）", extra={"error": str(e)})
|
||
return {}
|
||
|
||
def _replace_variables(
|
||
self,
|
||
text: str,
|
||
values: Dict[str, Any],
|
||
definitions: List[Dict[str, Any]]
|
||
) -> str:
|
||
"""替换文本中的变量
|
||
|
||
Args:
|
||
text: 原始文本
|
||
values: 变量值
|
||
definitions: 变量定义
|
||
|
||
Returns:
|
||
str: 替换后的文本
|
||
"""
|
||
result = text
|
||
|
||
# 创建变量定义映射
|
||
var_defs = {var["name"]: var for var in definitions}
|
||
|
||
for var_name, var_value in values.items():
|
||
# 检查变量是否在定义中
|
||
if var_name not in var_defs:
|
||
logger.warning(f"未定义的变量: {var_name}")
|
||
continue
|
||
|
||
# 替换变量（支持多种格式）
|
||
placeholders = [
|
||
f"{{{{{var_name}}}}}", # {{var_name}}
|
||
f"{{{var_name}}}", # {var_name}
|
||
f"${{{var_name}}}", # ${var_name}
|
||
]
|
||
|
||
for placeholder in placeholders:
|
||
if placeholder in result:
|
||
result = result.replace(placeholder, str(var_value))
|
||
|
||
return result
|
||
|
||
# ==================== 多模型对比试运行 ====================
|
||
|
||
async def run_compare(
|
||
self,
|
||
*,
|
||
agent_config: AgentConfig,
|
||
models: List[Dict[str, Any]],
|
||
message: str,
|
||
workspace_id: uuid.UUID,
|
||
conversation_id: Optional[str] = None,
|
||
user_id: Optional[str] = None,
|
||
variables: Optional[Dict[str, Any]] = None,
|
||
parallel: bool = True,
|
||
timeout: int = 60,
|
||
storage_type: Optional[str] = None,
|
||
user_rag_memory_id: Optional[str] = None,
|
||
web_search: bool = True,
|
||
memory: bool = True,
|
||
) -> Dict[str, Any]:
|
||
"""多模型对比试运行
|
||
|
||
Args:
|
||
agent_config: Agent 配置
|
||
models: 模型配置列表，每项包含 model_config, parameters, label, model_config_id
|
||
message: 用户消息
|
||
workspace_id: 工作空间ID
|
||
conversation_id: 会话ID
|
||
user_id: 用户ID
|
||
variables: 变量参数
|
||
parallel: 是否并行执行
|
||
timeout: 超时时间（秒）
|
||
|
||
Returns:
|
||
Dict: 对比结果
|
||
"""
|
||
logger.info(
|
||
f"多模型对比试运行",
|
||
extra={
|
||
"model_count": len(models),
|
||
"parallel": parallel
|
||
}
|
||
)
|
||
|
||
async def run_single_model(model_info):
|
||
"""运行单个模型"""
|
||
try:
|
||
start_time = time.time()
|
||
|
||
# 临时修改参数（不使用 deepcopy 避免 SQLAlchemy 会话问题）
|
||
original_params = agent_config.model_parameters
|
||
agent_config.model_parameters = model_info["parameters"]
|
||
|
||
# 使用模型自己的 conversation_id，如果没有则使用全局的
|
||
model_conversation_id = model_info.get("conversation_id") or conversation_id
|
||
try:
|
||
result = await asyncio.wait_for(
|
||
self.run(
|
||
agent_config=agent_config,
|
||
model_config=model_info["model_config"],
|
||
message=message,
|
||
workspace_id=workspace_id,
|
||
conversation_id=model_conversation_id,
|
||
user_id=user_id,
|
||
variables=variables,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id,
|
||
web_search=web_search,
|
||
memory=memory
|
||
),
|
||
timeout=timeout
|
||
)
|
||
finally:
|
||
# 恢复原始参数
|
||
agent_config.model_parameters = original_params
|
||
|
||
elapsed = time.time() - start_time
|
||
usage = result.get("usage", {})
|
||
|
||
return {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_info["label"],
|
||
"conversation_id":result['conversation_id'],
|
||
"parameters_used": model_info["parameters"],
|
||
"message": result.get("message"),
|
||
"usage": usage,
|
||
"elapsed_time": elapsed,
|
||
"tokens_per_second": (
|
||
usage.get("completion_tokens", 0) / elapsed
|
||
if elapsed > 0 and usage.get("completion_tokens") else None
|
||
),
|
||
"cost_estimate": self._estimate_cost(usage, model_info["model_config"]),
|
||
"error": None
|
||
}
|
||
|
||
except asyncio.TimeoutError:
|
||
logger.warning(
|
||
f"模型运行超时",
|
||
extra={
|
||
"model_config_id": str(model_info["model_config_id"]),
|
||
"timeout": timeout
|
||
}
|
||
)
|
||
return {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"conversation_id": conversation_id,
|
||
"label": model_info["label"],
|
||
"parameters_used": model_info["parameters"],
|
||
"elapsed_time": timeout,
|
||
"error": f"执行超时（{timeout}秒）"
|
||
}
|
||
except Exception as e:
|
||
logger.error(
|
||
f"模型运行失败",
|
||
extra={
|
||
"model_config_id": str(model_info["model_config_id"]),
|
||
"error": str(e)
|
||
}
|
||
)
|
||
return {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_info["label"],
|
||
"conversation_id": conversation_id,
|
||
"parameters_used": model_info["parameters"],
|
||
"elapsed_time": 0,
|
||
"error": str(e)
|
||
}
|
||
|
||
# 执行所有模型（并行或串行）
|
||
if parallel:
|
||
logger.debug(f"并行执行 {len(models)} 个模型")
|
||
results = await asyncio.gather(
|
||
*[run_single_model(m) for m in models],
|
||
return_exceptions=False
|
||
)
|
||
else:
|
||
logger.debug(f"串行执行 {len(models)} 个模型")
|
||
results = []
|
||
for model_info in models:
|
||
result = await run_single_model(model_info)
|
||
results.append(result)
|
||
|
||
# 统计分析
|
||
successful = [r for r in results if not r.get("error")]
|
||
failed = [r for r in results if r.get("error")]
|
||
|
||
fastest = min(successful, key=lambda x: x["elapsed_time"]) if successful else None
|
||
cheapest = min(
|
||
successful,
|
||
key=lambda x: x.get("cost_estimate") or float("inf")
|
||
) if successful else None
|
||
|
||
logger.info(
|
||
f"多模型对比完成",
|
||
extra={
|
||
"successful": len(successful),
|
||
"failed": len(failed),
|
||
"total_time": sum(r.get("elapsed_time", 0) for r in results)
|
||
}
|
||
)
|
||
|
||
return {
|
||
"results": results,
|
||
"total_elapsed_time": sum(r.get("elapsed_time", 0) for r in results),
|
||
"successful_count": len(successful),
|
||
"failed_count": len(failed),
|
||
"fastest_model": fastest["label"] if fastest else None,
|
||
"cheapest_model": cheapest["label"] if cheapest else None
|
||
}
|
||
|
||
def _estimate_cost(self, usage: Dict[str, Any], model_config) -> Optional[float]:
|
||
"""估算成本
|
||
|
||
Args:
|
||
usage: Token 使用情况
|
||
model_config: 模型配置
|
||
|
||
Returns:
|
||
Optional[float]: 估算成本（美元）
|
||
"""
|
||
if not usage:
|
||
return None
|
||
|
||
prompt_tokens = usage.get("prompt_tokens", 0)
|
||
completion_tokens = usage.get("completion_tokens", 0)
|
||
|
||
# 简化成本估算：暂时返回 None
|
||
# TODO: 实现基于模型名称或配置的成本估算
|
||
# 需要从 ModelApiKey 获取实际的模型名称，或者在 ModelConfig 中添加 model 字段
|
||
return None
|
||
|
||
def _with_parameters(self, agent_config: AgentConfig, parameters: Dict[str, Any]) -> AgentConfig:
|
||
"""创建一个带有覆盖参数的 agent_config（浅拷贝，只修改 model_parameters）
|
||
|
||
Args:
|
||
agent_config: 原始 Agent 配置
|
||
parameters: 要覆盖的参数
|
||
|
||
Returns:
|
||
AgentConfig: 修改后的配置（注意：这是同一个对象，只是临时修改了 model_parameters）
|
||
"""
|
||
# 保存原始参数
|
||
original_params = agent_config.model_parameters
|
||
# 设置新参数
|
||
agent_config.model_parameters = parameters
|
||
return agent_config, original_params
|
||
|
||
async def run_compare_stream(
|
||
self,
|
||
*,
|
||
agent_config: AgentConfig,
|
||
models: List[Dict[str, Any]],
|
||
message: str,
|
||
workspace_id: uuid.UUID,
|
||
conversation_id: Optional[str] = None,
|
||
user_id: Optional[str] = None,
|
||
variables: Optional[Dict[str, Any]] = None,
|
||
storage_type: Optional[str] = None,
|
||
user_rag_memory_id: Optional[str] = None,
|
||
web_search: bool = True,
|
||
memory: bool = True,
|
||
parallel: bool = True,
|
||
timeout: int = 60
|
||
) -> AsyncGenerator[str, None]:
|
||
"""多模型对比试运行（流式返回）
|
||
|
||
支持并行或串行执行，通过 model_index 区分不同模型的事件
|
||
|
||
Args:
|
||
agent_config: Agent 配置
|
||
models: 模型配置列表
|
||
message: 用户消息
|
||
workspace_id: 工作空间ID
|
||
conversation_id: 会话ID
|
||
user_id: 用户ID
|
||
variables: 变量参数
|
||
parallel: 是否并行执行
|
||
timeout: 超时时间（秒）
|
||
|
||
Yields:
|
||
str: SSE 格式的事件数据
|
||
"""
|
||
logger.info(
|
||
f"多模型对比流式试运行",
|
||
extra={"model_count": len(models), "parallel": parallel}
|
||
)
|
||
|
||
# 确保有 conversation_id
|
||
if not conversation_id:
|
||
conversation_id = str(uuid.uuid4())
|
||
|
||
# 发送开始事件
|
||
yield self._format_sse_event("compare_start", {
|
||
"conversation_id": conversation_id,
|
||
"model_count": len(models),
|
||
"parallel": parallel,
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
results = []
|
||
|
||
if parallel:
|
||
# 并行执行所有模型
|
||
import asyncio
|
||
|
||
# 创建事件队列用于收集所有模型的事件
|
||
event_queue = asyncio.Queue()
|
||
|
||
async def run_single_model_stream(idx: int, model_info: Dict[str, Any]):
|
||
"""运行单个模型并将事件放入队列"""
|
||
model_label = model_info["label"]
|
||
model_config_id = str(model_info["model_config_id"])
|
||
# 使用模型自己的 conversation_id，如果没有则使用全局的
|
||
model_conversation_id = model_info.get("conversation_id") or conversation_id
|
||
|
||
try:
|
||
# 发送模型开始事件
|
||
await event_queue.put(self._format_sse_event("model_start", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"timestamp": time.time()
|
||
}))
|
||
|
||
start_time = time.time()
|
||
full_content = ""
|
||
|
||
# 临时修改参数（并行任务中安全）
|
||
original_params = agent_config.model_parameters
|
||
agent_config.model_parameters = model_info["parameters"]
|
||
|
||
try:
|
||
# 流式调用单个模型
|
||
async for event_str in self.run_stream(
|
||
agent_config=agent_config,
|
||
model_config=model_info["model_config"],
|
||
message=message,
|
||
workspace_id=workspace_id,
|
||
conversation_id=model_conversation_id,
|
||
user_id=user_id,
|
||
variables=variables,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id,
|
||
web_search=web_search,
|
||
memory=memory
|
||
):
|
||
# 解析原始事件
|
||
try:
|
||
lines = event_str.strip().split('\n')
|
||
event_type = None
|
||
event_data = None
|
||
|
||
for line in lines:
|
||
if line.startswith('event: '):
|
||
event_type = line[7:].strip()
|
||
elif line.startswith('data: '):
|
||
event_data = json.loads(line[6:])
|
||
|
||
# 从 start 事件中获取 conversation_id
|
||
if event_type == "start" and event_data:
|
||
returned_conv_id = event_data.get("conversation_id")
|
||
if returned_conv_id:
|
||
model_conversation_id = returned_conv_id
|
||
|
||
if event_type == "message" and event_data:
|
||
chunk = event_data.get("content", "")
|
||
full_content += chunk
|
||
|
||
# 转发消息块事件（带模型标识和 conversation_id）
|
||
await event_queue.put(self._format_sse_event("model_message", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"content": chunk
|
||
}))
|
||
except Exception as e:
|
||
logger.warning(f"解析流式事件失败: {e}")
|
||
finally:
|
||
# 恢复原始参数
|
||
agent_config.model_parameters = original_params
|
||
|
||
elapsed = time.time() - start_time
|
||
|
||
# 模型完成
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"parameters_used": model_info["parameters"],
|
||
"message": full_content,
|
||
"elapsed_time": elapsed,
|
||
"error": None
|
||
}
|
||
|
||
# 发送模型完成事件
|
||
await event_queue.put(self._format_sse_event("model_end", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"elapsed_time": elapsed,
|
||
"message_length": len(full_content),
|
||
"timestamp": time.time()
|
||
}))
|
||
|
||
return result
|
||
|
||
except asyncio.TimeoutError:
|
||
logger.warning(f"模型运行超时: {model_label}")
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"elapsed_time": timeout,
|
||
"error": f"执行超时（{timeout}秒）"
|
||
}
|
||
|
||
await event_queue.put(self._format_sse_event("model_error", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"error": result["error"],
|
||
"timestamp": time.time()
|
||
}))
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.error(f"模型运行失败: {model_label}, error: {e}")
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"elapsed_time": 0,
|
||
"error": str(e)
|
||
}
|
||
|
||
await event_queue.put(self._format_sse_event("model_error", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"error": str(e),
|
||
"timestamp": time.time()
|
||
}))
|
||
|
||
return result
|
||
|
||
# 启动所有模型的并行任务
|
||
tasks = [
|
||
asyncio.create_task(run_single_model_stream(idx, model_info))
|
||
for idx, model_info in enumerate(models)
|
||
]
|
||
|
||
# 持续从队列中取出事件并发送
|
||
completed_count = 0
|
||
while completed_count < len(models):
|
||
try:
|
||
# 等待事件或任务完成
|
||
event = await asyncio.wait_for(event_queue.get(), timeout=0.1)
|
||
yield event
|
||
except asyncio.TimeoutError:
|
||
# 检查是否有任务完成
|
||
for task in tasks:
|
||
if task.done() and task not in [t for t in tasks if hasattr(t, '_result_retrieved')]:
|
||
result = await task
|
||
results.append(result)
|
||
task._result_retrieved = True
|
||
completed_count += 1
|
||
continue
|
||
|
||
# 等待所有任务完成
|
||
all_results = await asyncio.gather(*tasks, return_exceptions=False)
|
||
results = [r for r in all_results if r not in results]
|
||
results.extend([r for r in all_results if r not in results])
|
||
|
||
# 清空队列中剩余的事件
|
||
while not event_queue.empty():
|
||
try:
|
||
event = event_queue.get_nowait()
|
||
yield event
|
||
except asyncio.QueueEmpty:
|
||
break
|
||
|
||
else:
|
||
# 串行执行每个模型
|
||
for idx, model_info in enumerate(models):
|
||
model_label = model_info["label"]
|
||
model_config_id = str(model_info["model_config_id"])
|
||
# 使用模型自己的 conversation_id，如果没有则使用全局的
|
||
model_conversation_id = model_info.get("conversation_id") or conversation_id
|
||
|
||
# 发送模型开始事件
|
||
yield self._format_sse_event("model_start", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"conversation_id": model_conversation_id,
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
try:
|
||
start_time = time.time()
|
||
full_content = ""
|
||
|
||
# 临时修改参数
|
||
original_params = agent_config.model_parameters
|
||
agent_config.model_parameters = model_info["parameters"]
|
||
|
||
try:
|
||
# 流式调用单个模型
|
||
async for event_str in self.run_stream(
|
||
agent_config=agent_config,
|
||
model_config=model_info["model_config"],
|
||
message=message,
|
||
workspace_id=workspace_id,
|
||
conversation_id=model_conversation_id,
|
||
user_id=user_id,
|
||
variables=variables,
|
||
storage_type=storage_type,
|
||
user_rag_memory_id=user_rag_memory_id,
|
||
web_search=web_search,
|
||
memory=memory
|
||
):
|
||
# 解析原始事件
|
||
try:
|
||
# SSE 格式: "event: xxx\ndata: {...}\n\n"
|
||
lines = event_str.strip().split('\n')
|
||
event_type = None
|
||
event_data = None
|
||
|
||
for line in lines:
|
||
if line.startswith('event: '):
|
||
event_type = line[7:].strip()
|
||
elif line.startswith('data: '):
|
||
event_data = json.loads(line[6:])
|
||
|
||
if event_type == "message" and event_data:
|
||
# 累积内容
|
||
chunk = event_data.get("content", "")
|
||
full_content += chunk
|
||
|
||
# 转发消息块事件（带模型标识）
|
||
yield self._format_sse_event("model_message", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"content": chunk
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.warning(f"解析流式事件失败: {e}")
|
||
finally:
|
||
# 恢复原始参数
|
||
agent_config.model_parameters = original_params
|
||
|
||
elapsed = time.time() - start_time
|
||
|
||
# 模型完成
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"parameters_used": model_info["parameters"],
|
||
"message": full_content,
|
||
"elapsed_time": elapsed,
|
||
"error": None
|
||
}
|
||
results.append(result)
|
||
|
||
# 发送模型完成事件
|
||
yield self._format_sse_event("model_end", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"elapsed_time": elapsed,
|
||
"message_length": len(full_content),
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
except asyncio.TimeoutError:
|
||
logger.warning(f"模型运行超时: {model_label}")
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"elapsed_time": timeout,
|
||
"error": f"执行超时（{timeout}秒）"
|
||
}
|
||
results.append(result)
|
||
|
||
# 发送模型错误事件
|
||
yield self._format_sse_event("model_error", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"error": result["error"],
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
except Exception as e:
|
||
logger.error(f"模型运行失败: {model_label}, error: {e}")
|
||
result = {
|
||
"model_config_id": model_info["model_config_id"],
|
||
"model_name": model_info["model_config"].name,
|
||
"label": model_label,
|
||
"elapsed_time": 0,
|
||
"error": str(e)
|
||
}
|
||
results.append(result)
|
||
|
||
# 发送模型错误事件
|
||
yield self._format_sse_event("model_error", {
|
||
"model_index": idx,
|
||
"model_config_id": model_config_id,
|
||
"label": model_label,
|
||
"error": str(e),
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
# 统计分析
|
||
successful = [r for r in results if not r.get("error")]
|
||
failed = [r for r in results if r.get("error")]
|
||
|
||
fastest = min(successful, key=lambda x: x["elapsed_time"]) if successful else None
|
||
|
||
# 发送对比完成事件
|
||
yield self._format_sse_event("compare_end", {
|
||
"conversation_id": conversation_id,
|
||
"total_elapsed_time": sum(r.get("elapsed_time", 0) for r in results),
|
||
"successful_count": len(successful),
|
||
"failed_count": len(failed),
|
||
"fastest_model": fastest["label"] if fastest else None,
|
||
"timestamp": time.time()
|
||
})
|
||
|
||
logger.info(
|
||
f"多模型对比流式完成",
|
||
extra={
|
||
"successful": len(successful),
|
||
"failed": len(failed)
|
||
}
|
||
)
|
||
|
||
|
||
async def draft_run(
|
||
db: Session,
|
||
*,
|
||
agent_config: AgentConfig,
|
||
model_config: ModelConfig,
|
||
message: str,
|
||
user_id: Optional[str] = None,
|
||
kb_ids: Optional[List[str]] = None,
|
||
similarity_threshold: float = 0.7,
|
||
top_k: int = 3
|
||
) -> Dict[str, Any]:
|
||
"""试运行 Agent（便捷函数）
|
||
|
||
Args:
|
||
db: 数据库会话
|
||
agent_config: Agent 配置
|
||
model_config: 模型配置
|
||
message: 用户消息
|
||
user_id: 用户ID
|
||
kb_ids: 知识库ID列表
|
||
similarity_threshold: 相似度阈值
|
||
top_k: 检索返回的文档数量
|
||
|
||
Returns:
|
||
Dict: 包含 AI 回复和元数据的字典
|
||
"""
|
||
service = DraftRunService(db)
|
||
return await service.run(
|
||
agent_config=agent_config,
|
||
model_config=model_config,
|
||
message=message,
|
||
user_id=user_id,
|
||
kb_ids=kb_ids,
|
||
similarity_threshold=similarity_threshold,
|
||
top_k=top_k
|
||
)
|
||
|