6 changed files with 540 additions and 36 deletions
@ -0,0 +1,138 @@ |
|||||
|
import json |
||||
|
import time |
||||
|
from typing import Dict, List, Optional, Callable, Any |
||||
|
from enum import Enum |
||||
|
from dataclasses import dataclass, asdict |
||||
|
from datetime import datetime |
||||
|
|
||||
|
|
||||
|
class MessageType(Enum): |
||||
|
"""消息类型枚举""" |
||||
|
SEARCH = "search" |
||||
|
THINK = "think" |
||||
|
ANSWER = "answer" |
||||
|
|
||||
|
|
||||
|
@dataclass |
||||
|
class Message: |
||||
|
"""消息数据结构""" |
||||
|
type: MessageType |
||||
|
content: str |
||||
|
timestamp: float |
||||
|
metadata: Optional[Dict[str, Any]] = None |
||||
|
|
||||
|
def to_dict(self) -> Dict[str, Any]: |
||||
|
"""转换为字典格式""" |
||||
|
return { |
||||
|
"type": self.type.value, |
||||
|
"content": self.content, |
||||
|
"timestamp": self.timestamp, |
||||
|
"metadata": self.metadata or {} |
||||
|
} |
||||
|
|
||||
|
def to_json(self) -> str: |
||||
|
"""转换为JSON字符串""" |
||||
|
return json.dumps(self.to_dict(), ensure_ascii=False) |
||||
|
|
||||
|
|
||||
|
class MessageStream: |
||||
|
"""消息流管理器""" |
||||
|
|
||||
|
def __init__(self): |
||||
|
self._messages: List[Message] = [] |
||||
|
self._callbacks: List[Callable[[Message], None]] = [] |
||||
|
self._enabled = True |
||||
|
|
||||
|
def add_callback(self, callback: Callable[[Message], None]): |
||||
|
"""添加消息回调函数""" |
||||
|
self._callbacks.append(callback) |
||||
|
|
||||
|
def remove_callback(self, callback: Callable[[Message], None]): |
||||
|
"""移除消息回调函数""" |
||||
|
if callback in self._callbacks: |
||||
|
self._callbacks.remove(callback) |
||||
|
|
||||
|
def enable(self): |
||||
|
"""启用消息流""" |
||||
|
self._enabled = True |
||||
|
|
||||
|
def disable(self): |
||||
|
"""禁用消息流""" |
||||
|
self._enabled = False |
||||
|
|
||||
|
def send_message(self, msg_type: MessageType, content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""发送消息""" |
||||
|
if not self._enabled: |
||||
|
return |
||||
|
|
||||
|
message = Message( |
||||
|
type=msg_type, |
||||
|
content=content, |
||||
|
timestamp=time.time(), |
||||
|
metadata=metadata |
||||
|
) |
||||
|
|
||||
|
self._messages.append(message) |
||||
|
|
||||
|
# 调用所有回调函数 |
||||
|
for callback in self._callbacks: |
||||
|
try: |
||||
|
callback(message) |
||||
|
except Exception as e: |
||||
|
print(f"Error in message callback: {e}") |
||||
|
|
||||
|
def send_search(self, content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""发送搜索消息""" |
||||
|
self.send_message(MessageType.SEARCH, content, metadata) |
||||
|
|
||||
|
def send_think(self, content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""发送思考消息""" |
||||
|
self.send_message(MessageType.THINK, content, metadata) |
||||
|
|
||||
|
def send_answer(self, content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""发送答案消息""" |
||||
|
self.send_message(MessageType.ANSWER, content, metadata) |
||||
|
|
||||
|
def get_messages(self) -> List[Message]: |
||||
|
"""获取所有消息""" |
||||
|
return self._messages.copy() |
||||
|
|
||||
|
def get_messages_by_type(self, msg_type: MessageType) -> List[Message]: |
||||
|
"""根据类型获取消息""" |
||||
|
return [msg for msg in self._messages if msg.type == msg_type] |
||||
|
|
||||
|
def clear_messages(self): |
||||
|
"""清空所有消息""" |
||||
|
self._messages.clear() |
||||
|
|
||||
|
def get_messages_as_dicts(self) -> List[Dict[str, Any]]: |
||||
|
"""获取所有消息的字典格式""" |
||||
|
return [msg.to_dict() for msg in self._messages] |
||||
|
|
||||
|
def get_messages_as_json(self) -> str: |
||||
|
"""获取所有消息的JSON格式""" |
||||
|
return json.dumps(self.get_messages_as_dicts(), ensure_ascii=False) |
||||
|
|
||||
|
|
||||
|
# 全局消息流实例 |
||||
|
message_stream = MessageStream() |
||||
|
|
||||
|
|
||||
|
def send_search(content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""全局搜索消息发送函数""" |
||||
|
message_stream.send_search(content, metadata) |
||||
|
|
||||
|
|
||||
|
def send_think(content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""全局思考消息发送函数""" |
||||
|
message_stream.send_think(content, metadata) |
||||
|
|
||||
|
|
||||
|
def send_answer(content: str, metadata: Optional[Dict[str, Any]] = None): |
||||
|
"""全局答案消息发送函数""" |
||||
|
message_stream.send_answer(content, metadata) |
||||
|
|
||||
|
|
||||
|
def get_message_stream() -> MessageStream: |
||||
|
"""获取全局消息流实例""" |
||||
|
return message_stream |
||||
@ -0,0 +1,197 @@ |
|||||
|
# 消息流系统 (Message Stream System) |
||||
|
|
||||
|
## 概述 |
||||
|
|
||||
|
DeepSearcher 的消息流系统是一个新的消息传输机制，用于替代原来的日志传输方式。该系统支持三种类型的消息：`search`、`think` 和 `answer`，并提供了灵活的消息管理和传输功能。 |
||||
|
|
||||
|
## 消息类型 |
||||
|
|
||||
|
### 1. Search 消息 |
||||
|
- **类型**: `search` |
||||
|
- **用途**: 表示搜索相关的操作和状态 |
||||
|
- **示例**: |
||||
|
- "在向量数据库中搜索相关信息..." |
||||
|
- "找到5个相关文档片段" |
||||
|
- "搜索人工智能的定义..." |
||||
|
|
||||
|
### 2. Think 消息 |
||||
|
- **类型**: `think` |
||||
|
- **用途**: 表示思考和推理过程 |
||||
|
- **示例**: |
||||
|
- "开始处理查询: 什么是人工智能?" |
||||
|
- "分析搜索结果..." |
||||
|
- "生成子查询: 人工智能的定义、历史、应用" |
||||
|
|
||||
|
### 3. Answer 消息 |
||||
|
- **类型**: `answer` |
||||
|
- **用途**: 表示最终答案和结果 |
||||
|
- **示例**: |
||||
|
- "==== FINAL ANSWER====" |
||||
|
- "人工智能是计算机科学的一个分支..." |
||||
|
|
||||
|
## 核心组件 |
||||
|
|
||||
|
### MessageStream 类 |
||||
|
|
||||
|
主要的消息流管理器，提供以下功能： |
||||
|
|
||||
|
```python |
||||
|
from deepsearcher.utils.message_stream import MessageStream |
||||
|
|
||||
|
# 创建消息流实例 |
||||
|
message_stream = MessageStream() |
||||
|
|
||||
|
# 发送消息 |
||||
|
message_stream.send_search("搜索内容...") |
||||
|
message_stream.send_think("思考内容...") |
||||
|
message_stream.send_answer("答案内容...") |
||||
|
|
||||
|
# 获取消息 |
||||
|
messages = message_stream.get_messages() |
||||
|
messages_dict = message_stream.get_messages_as_dicts() |
||||
|
messages_json = message_stream.get_messages_as_json() |
||||
|
``` |
||||
|
|
||||
|
### 全局函数 |
||||
|
|
||||
|
为了方便使用，提供了全局函数： |
||||
|
|
||||
|
```python |
||||
|
from deepsearcher.utils.message_stream import send_search, send_think, send_answer |
||||
|
|
||||
|
# 直接发送消息 |
||||
|
send_search("搜索内容...") |
||||
|
send_think("思考内容...") |
||||
|
send_answer("答案内容...") |
||||
|
``` |
||||
|
|
||||
|
## API 接口 |
||||
|
|
||||
|
### 1. 获取消息 |
||||
|
``` |
||||
|
GET /messages/ |
||||
|
``` |
||||
|
|
||||
|
返回所有消息的列表： |
||||
|
```json |
||||
|
{ |
||||
|
"messages": [ |
||||
|
{ |
||||
|
"type": "search", |
||||
|
"content": "搜索内容...", |
||||
|
"timestamp": 1755043653.9606102, |
||||
|
"metadata": {} |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
``` |
||||
|
|
||||
|
### 2. 清空消息 |
||||
|
``` |
||||
|
POST /clear-messages/ |
||||
|
``` |
||||
|
|
||||
|
清空所有消息并返回成功状态： |
||||
|
```json |
||||
|
{ |
||||
|
"message": "Messages cleared successfully" |
||||
|
} |
||||
|
``` |
||||
|
|
||||
|
### 3. 查询接口（已更新） |
||||
|
``` |
||||
|
GET /query/?original_query=<query>&max_iter=<iterations> |
||||
|
``` |
||||
|
|
||||
|
现在返回结果包含消息流： |
||||
|
```json |
||||
|
{ |
||||
|
"result": "最终答案...", |
||||
|
"messages": [ |
||||
|
{ |
||||
|
"type": "search", |
||||
|
"content": "搜索内容...", |
||||
|
"timestamp": 1755043653.9606102, |
||||
|
"metadata": {} |
||||
|
} |
||||
|
] |
||||
|
} |
||||
|
``` |
||||
|
|
||||
|
## 前端集成 |
||||
|
|
||||
|
前端界面现在包含一个消息流显示区域，实时显示处理过程中的各种消息： |
||||
|
|
||||
|
### CSS 样式 |
||||
|
- `.message-search`: 搜索消息样式（蓝色边框） |
||||
|
- `.message-think`: 思考消息样式（黄色边框） |
||||
|
- `.message-answer`: 答案消息样式（绿色边框） |
||||
|
|
||||
|
### JavaScript 功能 |
||||
|
- `displayMessages(messages)`: 显示消息流 |
||||
|
- 自动滚动到最新消息 |
||||
|
- 时间戳显示 |
||||
|
|
||||
|
## 使用示例 |
||||
|
|
||||
|
### 后端使用 |
||||
|
|
||||
|
```python |
||||
|
from deepsearcher.utils.message_stream import send_search, send_think, send_answer |
||||
|
|
||||
|
# 在搜索过程中发送消息 |
||||
|
send_think("开始处理查询...") |
||||
|
send_search("在数据库中搜索...") |
||||
|
send_search("找到相关文档...") |
||||
|
send_think("分析结果...") |
||||
|
send_answer("最终答案...") |
||||
|
``` |
||||
|
|
||||
|
### 前端使用 |
||||
|
|
||||
|
```javascript |
||||
|
// 获取消息 |
||||
|
const response = await fetch('/query/?original_query=test&max_iter=3'); |
||||
|
const data = await response.json(); |
||||
|
|
||||
|
// 显示消息流 |
||||
|
if (data.messages && data.messages.length > 0) { |
||||
|
displayMessages(data.messages); |
||||
|
} |
||||
|
``` |
||||
|
|
||||
|
## 优势 |
||||
|
|
||||
|
1. **结构化数据**: 消息包含类型、内容、时间戳和元数据 |
||||
|
2. **类型安全**: 使用枚举确保消息类型的一致性 |
||||
|
3. **灵活传输**: 支持多种输出格式（字典、JSON） |
||||
|
4. **实时显示**: 前端可以实时显示处理过程 |
||||
|
5. **易于扩展**: 可以轻松添加新的消息类型和功能 |
||||
|
|
||||
|
## 迁移说明 |
||||
|
|
||||
|
从原来的日志系统迁移到新的消息流系统： |
||||
|
|
||||
|
### 原来的代码 |
||||
|
```python |
||||
|
log.color_print(f"<search> Search [{query}] in [{collection}]... </search>\n") |
||||
|
log.color_print(f"<think> Summarize answer... </think>\n") |
||||
|
log.color_print(f"<answer> Final answer... </answer>\n") |
||||
|
``` |
||||
|
|
||||
|
### 新的代码 |
||||
|
```python |
||||
|
send_search(f"Search [{query}] in [{collection}]...") |
||||
|
send_think("Summarize answer...") |
||||
|
send_answer("Final answer...") |
||||
|
``` |
||||
|
|
||||
|
## 测试 |
||||
|
|
||||
|
运行测试脚本验证系统功能： |
||||
|
|
||||
|
```bash |
||||
|
python test_message_stream.py |
||||
|
``` |
||||
|
|
||||
|
这将测试消息流的基本功能，包括消息发送、获取和格式化。 |
||||
@ -0,0 +1,31 @@ |
|||||
|
#!/usr/bin/env python3 |
||||
|
""" |
||||
|
测试迭代逻辑的脚本 |
||||
|
""" |
||||
|
|
||||
|
def test_iteration_logic(): |
||||
|
"""测试迭代逻辑""" |
||||
|
max_iter = 3 |
||||
|
|
||||
|
print(f"测试最大迭代次数: {max_iter}") |
||||
|
print("-" * 40) |
||||
|
|
||||
|
for it in range(max_iter): |
||||
|
print(f">> Iteration: {it + 1}") |
||||
|
|
||||
|
# 模拟搜索任务 |
||||
|
print(" 执行搜索任务...") |
||||
|
|
||||
|
# 检查是否达到最大迭代次数 |
||||
|
if it + 1 < max_iter: |
||||
|
print(" 反思并生成新的子查询...") |
||||
|
print(" 准备下一次迭代") |
||||
|
else: |
||||
|
print(" 达到最大迭代次数，退出") |
||||
|
break |
||||
|
|
||||
|
print("-" * 40) |
||||
|
print("测试完成!") |
||||
|
|
||||
|
if __name__ == "__main__": |
||||
|
test_iteration_logic() |
||||
Loading…
Reference in new issue