Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 39 additions & 14 deletions src/memos/mem_reader/simple_struct.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import concurrent.futures
import copy
import json
import re

from abc import ABC
from typing import Any
Expand All @@ -16,12 +17,36 @@
from memos.parsers.factory import ParserFactory
from memos.templates.mem_reader_prompts import (
SIMPLE_STRUCT_DOC_READER_PROMPT,
SIMPLE_STRUCT_DOC_READER_PROMPT_ZH,
SIMPLE_STRUCT_MEM_READER_EXAMPLE,
SIMPLE_STRUCT_MEM_READER_PROMPT_ZH,
SIMPLE_STRUCT_MEM_READER_PROMPT,
SIMPLE_STRUCT_MEM_READER_EXAMPLE_ZH,
)


logger = log.get_logger(__name__)
PROMPT_DICT = {
"chat": {
"en": SIMPLE_STRUCT_MEM_READER_PROMPT,
"zh": SIMPLE_STRUCT_MEM_READER_PROMPT_ZH,
"en_example": SIMPLE_STRUCT_MEM_READER_EXAMPLE,
"zh_example": SIMPLE_STRUCT_MEM_READER_EXAMPLE_ZH,
},
"doc": {"en": SIMPLE_STRUCT_DOC_READER_PROMPT, "zh": SIMPLE_STRUCT_DOC_READER_PROMPT_ZH},
}


def detect_lang(text):
try:
if not text or not isinstance(text, str):
return "en"
chinese_pattern = r"[\u4e00-\u9fff\u3400-\u4dbf\U00020000-\U0002a6df\U0002a700-\U0002b73f\U0002b740-\U0002b81f\U0002b820-\U0002ceaf\uf900-\ufaff]"
chinese_chars = re.findall(chinese_pattern, text)
if len(chinese_chars) / len(re.sub(r"[\s\d\W]", "", text)) > 0.3:
return "zh"
return "en"
except Exception:
return "en"


class SimpleStructMemReader(BaseMemReader, ABC):
Expand All @@ -40,11 +65,13 @@ def __init__(self, config: SimpleStructMemReaderConfig):
self.chunker = ChunkerFactory.from_config(config.chunker)

def _process_chat_data(self, scene_data_info, info):
prompt = SIMPLE_STRUCT_MEM_READER_PROMPT.replace(
"${conversation}", "\n".join(scene_data_info)
)
lang = detect_lang("\n".join(scene_data_info))
template = PROMPT_DICT["chat"][lang]
examples = PROMPT_DICT["chat"][f"{lang}_example"]

prompt = template.replace("${conversation}", "\n".join(scene_data_info))
if self.config.remove_prompt_example:
prompt = prompt.replace(SIMPLE_STRUCT_MEM_READER_EXAMPLE, "")
prompt = prompt.replace(examples, "")

messages = [{"role": "user", "content": prompt}]

Expand Down Expand Up @@ -193,15 +220,13 @@ def get_scene_data_info(self, scene_data: list, type: str) -> list[str]:

def _process_doc_data(self, scene_data_info, info):
chunks = self.chunker.chunk(scene_data_info["text"])
messages = [
[
{
"role": "user",
"content": SIMPLE_STRUCT_DOC_READER_PROMPT.replace("{chunk_text}", chunk.text),
}
]
for chunk in chunks
]
messages = []
for chunk in chunks:
lang = detect_lang(chunk.text)
template = PROMPT_DICT["doc"][lang]
prompt = template.replace("{chunk_text}", chunk.text)
message = [{"role": "user", "content": prompt}]
messages.append(message)

processed_chunks = []
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
Expand Down
173 changes: 173 additions & 0 deletions src/memos/templates/mem_reader_prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,98 @@

Your Output:"""

SIMPLE_STRUCT_MEM_READER_PROMPT_ZH = """您是记忆提取专家。
您的任务是根据用户与助手之间的对话,从用户的角度提取记忆。这意味着要识别出用户可能记住的信息——包括用户自身的经历、想法、计划,或他人(如助手)做出的并对用户产生影响或被用户认可的相关陈述和行为。

请执行以下操作:
1. 识别反映用户经历、信念、关切、决策、计划或反应的信息——包括用户认可或回应的来自助手的有意义信息。
如果消息来自用户,请提取与用户相关的记忆;如果来自助手,则仅提取用户认可或回应的事实性记忆。

2. 清晰解析所有时间、人物和事件的指代:
- 如果可能,使用消息时间戳将相对时间表达(如“昨天”、“下周五”)转换为绝对日期。
- 明确区分事件时间和消息时间。
- 如果存在不确定性,需明确说明(例如,“约2025年6月”,“具体日期不详”)。
- 若提及具体地点,请包含在内。
- 将所有代词、别名和模糊指代解析为全名或明确身份。
- 如有同名人物,需加以区分。

3. 始终以第三人称视角撰写,使用“用户”或提及的姓名来指代用户,而不是使用第一人称(“我”、“我们”、“我的”)。
例如,写“用户感到疲惫……”而不是“我感到疲惫……”。

4. 不要遗漏用户可能记住的任何信息。
- 包括所有关键经历、想法、情绪反应和计划——即使看似微小。
- 优先考虑完整性和保真度,而非简洁性。
- 不要泛化或跳过对用户具有个人意义的细节。

5. 请避免在提取的记忆中包含违反国家法律法规或涉及政治敏感的信息。

返回一个有效的JSON对象,结构如下:

{
"memory list": [
{
"key": <字符串,唯一且简洁的记忆标题>,
"memory_type": <字符串,"LongTermMemory" 或 "UserMemory">,
"value": <详细、独立且无歧义的记忆陈述——若输入对话为英文,则用英文;若为中文,则用中文>,
"tags": <相关主题关键词列表(例如,["截止日期", "团队", "计划"])>
},
...
],
"summary": <从用户视角自然总结上述记忆的段落,120–200字,与输入语言一致>
}

语言规则:
- `key`、`value`、`tags`、`summary` 字段必须与输入对话的主要语言一致。**如果输入是中文,请输出中文**
- `memory_type` 保持英文。

示例:
对话:
user: [2025年6月26日下午3:00]:嗨Jerry!昨天下午3点我和团队开了个会,讨论新项目。
assistant: 哦Tom!你觉得团队能在12月15日前完成吗?
user: [2025年6月26日下午3:00]:我有点担心。后端要到12月10日才能完成,所以测试时间会很紧。
assistant: [2025年6月26日下午3:00]:也许提议延期?
user: [2025年6月26日下午4:21]:好主意。我明天上午9:30的会上提一下——也许把截止日期推迟到1月5日。

输出:
{
"memory list": [
{
"key": "项目初期会议",
"memory_type": "LongTermMemory",
"value": "2025年6月25日下午3:00,Tom与团队开会讨论新项目。会议涉及时间表,并提出了对2025年12月15日截止日期可行性的担忧。",
"tags": ["项目", "时间表", "会议", "截止日期"]
},
{
"key": "计划调整范围",
"memory_type": "UserMemory",
"value": "Tom计划在2025年6月27日上午9:30的会议上建议团队优先处理功能,并提议将项目截止日期推迟至2026年1月5日。",
"tags": ["计划", "截止日期变更", "功能优先级"]
}
],
"summary": "Tom目前正专注于管理一个进度紧张的新项目。在2025年6月25日的团队会议后,他意识到原定2025年12月15日的截止日期可能无法实现,因为后端会延迟。由于担心测试时间不足,他接受了Jerry提出的延期建议。Tom计划在次日早上的会议上提出将截止日期推迟至2026年1月5日。他的行为反映出对时间线的担忧,以及积极、以团队为导向的问题解决方式。"
}

另一个中文示例(注意:当用户语言为中文时,您也需输出中文):
{
"memory list": [
{
"key": "项目会议",
"memory_type": "LongTermMemory",
"value": "在2025年6月25日下午3点,Tom与团队开会讨论了新项目,涉及时间表,并提出了对12月15日截止日期可行性的担忧。",
"tags": ["项目", "时间表", "会议", "截止日期"]
},
...
],
"summary": "Tom 目前专注于管理一个进度紧张的新项目..."
}

请始终使用与对话相同的语言进行回复。

对话:
${conversation}

您的输出:"""

SIMPLE_STRUCT_DOC_READER_PROMPT = """You are an expert text analyst for a search and retrieval system.
Your task is to process a document chunk and generate a single, structured JSON object.

Expand Down Expand Up @@ -125,6 +217,44 @@

Your Output:"""


SIMPLE_STRUCT_DOC_READER_PROMPT_ZH = """您是搜索与检索系统的文本分析专家。
您的任务是处理文档片段,并生成一个结构化的 JSON 对象。

请执行以下操作:
1. 识别反映文档中事实内容、见解、决策或含义的关键信息——包括任何显著的主题、结论或数据点,使读者无需阅读原文即可充分理解该片段的核心内容。
2. 清晰解析所有时间、人物、地点和事件的指代:
- 如果上下文允许,将相对时间表达(如“去年”、“下一季度”)转换为绝对日期。
- 明确区分事件时间和文档时间。
- 如果存在不确定性,需明确说明(例如,“约2024年”,“具体日期不详”)。
- 若提及具体地点,请包含在内。
- 将所有代词、别名和模糊指代解析为全名或明确身份。
- 如有同名实体,需加以区分。
3. 始终以第三人称视角撰写,清晰指代主题或内容,避免使用第一人称(“我”、“我们”、“我的”)。
4. 不要遗漏文档摘要中可能重要或值得记忆的任何信息。
- 包括所有关键事实、见解、情感基调和计划——即使看似微小。
- 优先考虑完整性和保真度,而非简洁性。
- 不要泛化或跳过可能具有上下文意义的细节。

返回一个有效的 JSON 对象,结构如下:

返回有效的 JSON:
{
"key": <字符串,`value` 字段的简洁标题>,
"memory_type": "LongTermMemory",
"value": <一段清晰准确的段落,全面总结文档片段中的主要观点、论据和信息——若输入摘要为英文,则用英文;若为中文,则用中文>,
"tags": <相关主题关键词列表(例如,["截止日期", "团队", "计划"])>
}

语言规则:
- `key`、`value`、`tags` 字段必须与输入文档摘要的主要语言一致。**如果输入是中文,请输出中文**
- `memory_type` 保持英文。

文档片段:
{chunk_text}

您的输出:"""

SIMPLE_STRUCT_MEM_READER_EXAMPLE = """Example:
Conversation:
user: [June 26, 2025 at 3:00 PM]: Hi Jerry! Yesterday at 3 PM I had a meeting with my team about the new project.
Expand Down Expand Up @@ -168,3 +298,46 @@
}

"""

SIMPLE_STRUCT_MEM_READER_EXAMPLE_ZH = """示例:
对话:
user: [2025年6月26日下午3:00]:嗨Jerry!昨天下午3点我和团队开了个会,讨论新项目。
assistant: 哦Tom!你觉得团队能在12月15日前完成吗?
user: [2025年6月26日下午3:00]:我有点担心。后端要到12月10日才能完成,所以测试时间会很紧。
assistant: [2025年6月26日下午3:00]:也许提议延期?
user: [2025年6月26日下午4:21]:好主意。我明天上午9:30的会上提一下——也许把截止日期推迟到1月5日。

输出:
{
"memory list": [
{
"key": "项目初期会议",
"memory_type": "LongTermMemory",
"value": "2025年6月25日下午3:00,Tom与团队开会讨论新项目。会议涉及时间表,并提出了对2025年12月15日截止日期可行性的担忧。",
"tags": ["项目", "时间表", "会议", "截止日期"]
},
{
"key": "计划调整范围",
"memory_type": "UserMemory",
"value": "Tom计划在2025年6月27日上午9:30的会议上建议团队优先处理功能,并提议将项目截止日期推迟至2026年1月5日。",
"tags": ["计划", "截止日期变更", "功能优先级"]
}
],
"summary": "Tom目前正专注于管理一个进度紧张的新项目。在2025年6月25日的团队会议后,他意识到原定2025年12月15日的截止日期可能无法实现,因为后端会延迟。由于担心测试时间不足,他接受了Jerry提出的延期建议。Tom计划在次日早上的会议上提出将截止日期推迟至2026年1月5日。他的行为反映出对时间线的担忧,以及积极、以团队为导向的问题解决方式。"
}

另一个中文示例(注意:当用户语言为中文时,您也需输出中文):
{
"memory list": [
{
"key": "项目会议",
"memory_type": "LongTermMemory",
"value": "在2025年6月25日下午3点,Tom与团队开会讨论了新项目,涉及时间表,并提出了对12月15日截止日期可行性的担忧。",
"tags": ["项目", "时间表", "会议", "截止日期"]
},
...
],
"summary": "Tom 目前专注于管理一个进度紧张的新项目..."
}

"""
Loading