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