227 lines
7.9 KiB
Python
227 lines
7.9 KiB
Python
"""聊天服务"""
|
||
|
||
import asyncio
|
||
import json
|
||
from datetime import datetime
|
||
from typing import AsyncIterator, TYPE_CHECKING, List, Dict, Any
|
||
from ..core.imports import GraphMemoryClient
|
||
from ..models.message import ToolCall, ToolResult
|
||
from .tool_service import ToolService
|
||
|
||
if TYPE_CHECKING:
|
||
from ..core.imports import Neo4jGraph
|
||
|
||
|
||
class ChatService:
|
||
"""聊天业务服务"""
|
||
|
||
def __init__(
|
||
self,
|
||
graph: "Neo4jGraph",
|
||
client: GraphMemoryClient,
|
||
tool_service: ToolService
|
||
):
|
||
self._graph = graph
|
||
self._client = client
|
||
self._tool_service = tool_service
|
||
self._messages: List[Dict[str, Any]] = []
|
||
|
||
async def send_message(self, user_input: str) -> AsyncIterator[dict]:
|
||
"""发送消息并流式返回事件"""
|
||
# 1. 发送用户消息事件
|
||
yield {
|
||
"type": "user_message",
|
||
"content": user_input
|
||
}
|
||
|
||
try:
|
||
# 2. 第一次API调用
|
||
accumulated_content = ""
|
||
tool_calls_data = []
|
||
|
||
# 流式处理响应
|
||
async for chunk in self._call_api_stream_async(user_input):
|
||
if chunk.get("content_delta"):
|
||
accumulated_content += chunk["content_delta"]
|
||
yield {
|
||
"type": "content_delta",
|
||
"content": accumulated_content
|
||
}
|
||
|
||
if chunk.get("tool_calls"):
|
||
tool_calls_data = chunk["tool_calls"]
|
||
|
||
# 3. 如果有工具调用，执行并继续调用API
|
||
tool_calls = None
|
||
tool_results = None
|
||
|
||
if tool_calls_data:
|
||
tool_calls = []
|
||
tool_results = []
|
||
|
||
# 执行所有工具
|
||
for tool_call_data in tool_calls_data:
|
||
tool_call = ToolCall(
|
||
id=tool_call_data["id"],
|
||
name=tool_call_data["function"]["name"],
|
||
arguments=tool_call_data["function"]["arguments"]
|
||
)
|
||
tool_calls.append(tool_call)
|
||
|
||
yield {
|
||
"type": "tool_call",
|
||
"tool_call": tool_call
|
||
}
|
||
|
||
result = await self._tool_service.execute(tool_call)
|
||
tool_results.append(result)
|
||
|
||
log_entry = ToolService._create_log_entry(tool_call, result)
|
||
yield {
|
||
"type": "tool_result",
|
||
"tool_result": result,
|
||
"log_entry": log_entry
|
||
}
|
||
|
||
# 构建工具结果消息
|
||
tool_messages = []
|
||
for tc, tr in zip(tool_calls, tool_results):
|
||
tool_messages.append({
|
||
"role": "tool",
|
||
"tool_call_id": tc.id,
|
||
"content": tr.content
|
||
})
|
||
|
||
# 构建assistant消息（包含tool_calls）
|
||
assistant_message = {
|
||
"role": "assistant",
|
||
"content": accumulated_content,
|
||
"tool_calls": [
|
||
{
|
||
"id": tc.id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": tc.name,
|
||
"arguments": tc.arguments
|
||
}
|
||
} for tc in tool_calls
|
||
]
|
||
}
|
||
|
||
# 第二次API调用，传入工具结果
|
||
final_content = ""
|
||
async for chunk in self._call_api_stream_with_tools(
|
||
user_input,
|
||
assistant_message,
|
||
tool_messages
|
||
):
|
||
if chunk.get("content_delta"):
|
||
final_content += chunk["content_delta"]
|
||
yield {
|
||
"type": "content_delta",
|
||
"content": final_content
|
||
}
|
||
|
||
accumulated_content = final_content
|
||
|
||
# 4. 返回最终回复
|
||
yield {
|
||
"type": "assistant_message",
|
||
"content": accumulated_content,
|
||
"tool_calls": tool_calls,
|
||
"tool_results": tool_results
|
||
}
|
||
|
||
except Exception as e:
|
||
yield {
|
||
"type": "error",
|
||
"error": str(e)
|
||
}
|
||
|
||
async def _call_api_stream_async(self, message: str) -> AsyncIterator[dict]:
|
||
"""异步流式调用 API"""
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def process_stream():
|
||
stream = self._client.send_message_stream(message)
|
||
tool_calls_accumulated = []
|
||
|
||
for chunk in stream:
|
||
delta = chunk.choices[0].delta
|
||
|
||
if delta.content:
|
||
yield {"content_delta": delta.content}
|
||
|
||
if delta.tool_calls:
|
||
for tc in delta.tool_calls:
|
||
if tc.index >= len(tool_calls_accumulated):
|
||
tool_calls_accumulated.append({
|
||
"id": tc.id,
|
||
"type": "function",
|
||
"function": {
|
||
"name": "",
|
||
"arguments": ""
|
||
}
|
||
})
|
||
|
||
if tc.function:
|
||
if tc.function.name:
|
||
tool_calls_accumulated[tc.index]["function"]["name"] = tc.function.name
|
||
if tc.function.arguments:
|
||
tool_calls_accumulated[tc.index]["function"]["arguments"] += tc.function.arguments
|
||
|
||
if tool_calls_accumulated:
|
||
yield {"tool_calls": tool_calls_accumulated}
|
||
|
||
for result in await loop.run_in_executor(None, lambda: list(process_stream())):
|
||
yield result
|
||
|
||
async def _call_api_stream_with_tools(
|
||
self,
|
||
user_input: str,
|
||
assistant_message: dict,
|
||
tool_messages: list
|
||
) -> AsyncIterator[dict]:
|
||
"""带工具结果的流式调用"""
|
||
loop = asyncio.get_event_loop()
|
||
|
||
def process_stream():
|
||
# 构建完整的消息列表
|
||
messages = [
|
||
{"role": "system", "content": self._client.system_prompt},
|
||
{"role": "user", "content": user_input},
|
||
assistant_message
|
||
]
|
||
messages.extend(tool_messages)
|
||
|
||
# 调用API
|
||
response = self._client.client.chat.completions.create(
|
||
model="deepseek-chat",
|
||
messages=messages,
|
||
tools=self._client.tools,
|
||
tool_choice="auto",
|
||
stream=True
|
||
)
|
||
|
||
for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
if delta.content:
|
||
yield {"content_delta": delta.content}
|
||
|
||
# 处理可能的工具调用
|
||
if delta.tool_calls:
|
||
# 如果还有工具调用，说明AI想继续调用工具
|
||
# 但我们限制只调用一次，所以忽略
|
||
pass
|
||
|
||
for result in await loop.run_in_executor(None, lambda: list(process_stream())):
|
||
yield result
|
||
|
||
def clear_history(self) -> None:
|
||
"""清空消息历史"""
|
||
self._messages.clear()
|
||
|
||
def get_history(self) -> list[dict]:
|
||
"""获取消息历史"""
|
||
return self._messages.copy()
|