From b06ff989eaff10417122caaa6b49c5a5d614c07e Mon Sep 17 00:00:00 2001 From: bittergreen Date: Fri, 24 Jul 2026 10:40:55 +0800 Subject: [PATCH] fix: Add dedicated document LLM, improve Markdown detection and context-aware structured extraction --- src/memos/api/config.py | 26 +++- src/memos/configs/mem_reader.py | 7 +- src/memos/mem_reader/multi_modal_struct.py | 9 +- .../read_multi_modal/file_content_parser.py | 50 ++++++- .../read_multi_modal/multi_modal_parser.py | 7 +- src/memos/templates/mem_reader_prompts.py | 46 +++++-- tests/api/test_llm_provider_config.py | 51 +++++++ tests/mem_reader/test_document_parser_llm.py | 128 ++++++++++++++++++ 8 files changed, 303 insertions(+), 21 deletions(-) create mode 100644 tests/mem_reader/test_document_parser_llm.py diff --git a/src/memos/api/config.py b/src/memos/api/config.py index b3e45dfb2..43ad586dd 100644 --- a/src/memos/api/config.py +++ b/src/memos/api/config.py @@ -400,7 +400,7 @@ def get_activation_config() -> dict[str, Any]: @staticmethod def get_memreader_config() -> dict[str, Any]: - """Get MemReader configuration for chat/doc extraction (fine-tuned 0.6B model). + """Get the main MemReader configuration for text/chat extraction. When MEMREADER_GENERAL_MODEL is configured (i.e. a separate stable LLM exists), the backup client is automatically enabled so that primary failures (self-deployed @@ -444,7 +444,7 @@ def get_qwen_llm_config() -> dict[str, Any] | None: @staticmethod def get_memreader_general_llm_config() -> dict[str, Any]: - """Get general LLM configuration for non-chat/doc tasks. + """Get general LLM configuration for non-primary extraction tasks. Used for: hallucination filter, memory rewrite, memory merge, tool trajectory extraction, skill memory extraction. @@ -478,6 +478,24 @@ def get_image_parser_llm_config() -> dict[str, Any]: # Fallback to general_llm config (which itself falls back to OpenAI) return APIConfig.get_memreader_general_llm_config() + @staticmethod + def get_document_parser_llm_config() -> dict[str, Any] | None: + """Get the dedicated LLM configuration for document extraction. + + The provider endpoint and credentials are selected from QWEN_* or + OPENAI_* according to DOCUMENT_PARSER_MODEL. + + Fallback chain: DOCUMENT_PARSER_MODEL -> MEMREADER_GENERAL_MODEL. + Returns None when neither model is configured. + """ + document_model = os.getenv("DOCUMENT_PARSER_MODEL") or os.getenv("MEMREADER_GENERAL_MODEL") + if not document_model: + return None + return APIConfig._build_provider_llm_config( + document_model, + temperature=0.8, + ) + @staticmethod def get_preference_extractor_llm_config() -> dict[str, Any]: """Get LLM configuration for preference extraction. @@ -994,6 +1012,8 @@ def get_product_default_config() -> dict[str, Any]: "general_llm": APIConfig.get_memreader_general_llm_config(), # Image parser LLM (requires vision model) "image_parser_llm": APIConfig.get_image_parser_llm_config(), + # Dedicated LLM for document chunk extraction + "document_parser_llm": APIConfig.get_document_parser_llm_config(), # Preference extractor LLM. Reader falls back to general_llm when unset. "preference_extractor_llm": APIConfig.get_preference_extractor_llm_config() if os.getenv("PREFERENCE_EXTRACTOR_MODEL") @@ -1127,6 +1147,8 @@ def create_user_config(user_name: str, user_id: str) -> tuple["MOSConfig", "Gene "general_llm": APIConfig.get_memreader_general_llm_config(), # Image parser LLM (requires vision model) "image_parser_llm": APIConfig.get_image_parser_llm_config(), + # Dedicated LLM for document chunk extraction + "document_parser_llm": APIConfig.get_document_parser_llm_config(), # Preference extractor LLM. Reader falls back to general_llm when unset. "preference_extractor_llm": APIConfig.get_preference_extractor_llm_config() if os.getenv("PREFERENCE_EXTRACTOR_MODEL") diff --git a/src/memos/configs/mem_reader.py b/src/memos/configs/mem_reader.py index c7ee93ad6..5bbcfa0c0 100644 --- a/src/memos/configs/mem_reader.py +++ b/src/memos/configs/mem_reader.py @@ -25,7 +25,8 @@ def parse_datetime(cls, value): return value llm: LLMConfigFactory = Field( - ..., description="LLM configuration for chat/doc memory extraction (fine-tuned model)" + ..., + description="Main LLM configuration for standard text/chat memory extraction", ) general_llm: LLMConfigFactory | None = Field( default=None, @@ -36,6 +37,10 @@ def parse_datetime(cls, value): default=None, description="Vision LLM for image parsing. Falls back to general_llm if not set.", ) + document_parser_llm: LLMConfigFactory | None = Field( + default=None, + description="Dedicated LLM for document content extraction", + ) preference_extractor_llm: LLMConfigFactory | None = Field( default=None, description="LLM for preference extraction. Falls back to general_llm if not set.", diff --git a/src/memos/mem_reader/multi_modal_struct.py b/src/memos/mem_reader/multi_modal_struct.py index 8a7bc2da1..c46331f1d 100644 --- a/src/memos/mem_reader/multi_modal_struct.py +++ b/src/memos/mem_reader/multi_modal_struct.py @@ -71,12 +71,19 @@ def __init__(self, config: MultiModalStructMemReaderConfig): if config.image_parser_llm is not None else self.general_llm ) + # Document parser LLM (separate from the fine-tuned chat extractor) + self.document_parser_llm = ( + LLMFactory.from_config(config.document_parser_llm) + if config.document_parser_llm is not None + else None + ) # Initialize MultiModalParser for routing to different parsers - # Pass image_parser_llm for image parsing + # Pass dedicated image/document LLMs to their corresponding parsers self.multi_modal_parser = MultiModalParser( embedder=self.embedder, llm=self.llm, image_parser_llm=self.image_parser_llm, + document_parser_llm=self.document_parser_llm, parser=None, direct_markdown_hostnames=direct_markdown_hostnames, ) diff --git a/src/memos/mem_reader/read_multi_modal/file_content_parser.py b/src/memos/mem_reader/read_multi_modal/file_content_parser.py index 0629eda97..a6aba3aa0 100644 --- a/src/memos/mem_reader/read_multi_modal/file_content_parser.py +++ b/src/memos/mem_reader/read_multi_modal/file_content_parser.py @@ -100,7 +100,7 @@ def _get_doc_llm_response( def _handle_url(self, url_str: str, filename: str) -> tuple[str, str | None, bool]: """Download and parse file from URL.""" try: - from urllib.parse import urlparse + from urllib.parse import unquote, urlparse import requests @@ -118,7 +118,13 @@ def _handle_url(self, url_str: str, filename: str) -> tuple[str, str | None, boo return response.text, None, True file_ext = os.path.splitext(filename)[1].lower() - if file_ext in [".md", ".markdown", ".txt"] or self._is_oss_md(url_str): + url_path_ext = os.path.splitext(unquote(parsed_url.path))[1].lower() + markdown_extensions = {".md", ".markdown", ".txt"} + if ( + file_ext in markdown_extensions + or url_path_ext in markdown_extensions + or self._is_oss_md(url_str) + ): return response.text, None, True with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=file_ext) as temp_file: temp_file.write(response.content) @@ -333,7 +339,7 @@ def __init__( Args: embedder: Embedder for generating embeddings - llm: Optional LLM for fine mode processing + llm: Optional dedicated LLM for document extraction parser: Optional parser for parsing file contents direct_markdown_hostnames: List of hostnames that should return markdown directly without parsing. If None, reads from FILE_PARSER_DIRECT_MARKDOWN_HOSTNAMES @@ -739,6 +745,13 @@ def parse_fine( logger.info( f"[Chunker: FileContentParser] Extracted {len(headers)} headers from markdown" ) + document_context = self._build_markdown_document_context(headers, filename) + if document_context: + context_parts = [] + if message_text_context: + context_parts.append(f"Related message context:\n{message_text_context}") + context_parts.append(document_context) + message_text_context = "\n\n".join(context_parts) # Extract and process images from parsed_text if is_markdown and parsed_text and self.image_parser: @@ -1050,6 +1063,37 @@ def _extract_markdown_headers(self, text: str) -> dict[int, dict]: logger.info(f"[Chunker: FileContentParser] Extracted {len(headers)} headers from markdown") return headers + @staticmethod + def _build_markdown_document_context( + headers: dict[int, dict], + filename: str, + ) -> str | None: + """Build document-level context from a Markdown filename and H1 titles.""" + h1_titles = [] + seen_titles = set() + for header in headers.values(): + if header.get("level") != 1: + continue + title = str(header.get("title", "")).strip() + if not title or title in seen_titles: + continue + seen_titles.add(title) + h1_titles.append(title) + + if not filename and not h1_titles: + return None + + lines = [ + "Document-level context for disambiguation only. " + "Do not create memories from this context unless the current chunk supports them." + ] + if filename: + lines.append(f"Filename: {filename}") + if h1_titles: + lines.append("Top-level Markdown headings:") + lines.extend(f"- {title}" for title in h1_titles) + return "\n".join(lines) + def _get_header_context( self, text: str, image_position: int, headers: dict[int, dict] ) -> list[str]: diff --git a/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py b/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py index a08aadc0c..bdcaee6ee 100644 --- a/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py +++ b/src/memos/mem_reader/read_multi_modal/multi_modal_parser.py @@ -38,6 +38,7 @@ def __init__( embedder: BaseEmbedder, llm: BaseLLM | None = None, image_parser_llm: BaseLLM | None = None, + document_parser_llm: BaseLLM | None = None, parser: Any | None = None, direct_markdown_hostnames: list[str] | None = None, ): @@ -46,9 +47,10 @@ def __init__( Args: embedder: Embedder for generating embeddings - llm: Optional LLM for fine mode processing (chat/doc extraction) + llm: Optional main LLM for standard text/chat extraction. image_parser_llm: Optional vision LLM for image parsing. Falls back to llm if not provided. + document_parser_llm: Optional dedicated LLM for document extraction. parser: Optional parser for parsing file contents direct_markdown_hostnames: List of hostnames that should return markdown directly without parsing. If None, reads from FILE_PARSER_DIRECT_MARKDOWN_HOSTNAMES @@ -58,6 +60,7 @@ def __init__( self.llm = llm # Image parser LLM (requires vision model), falls back to main llm self.image_parser_llm = image_parser_llm if image_parser_llm is not None else llm + self.document_parser_llm = document_parser_llm self.parser = parser # Initialize parsers for different message types @@ -71,7 +74,7 @@ def __init__( self.image_parser = ImageParser(embedder, self.image_parser_llm) self.file_content_parser = FileContentParser( embedder, - llm, + self.document_parser_llm, parser, direct_markdown_hostnames=direct_markdown_hostnames, image_parser=self.image_parser, diff --git a/src/memos/templates/mem_reader_prompts.py b/src/memos/templates/mem_reader_prompts.py index 63e4c1538..3055ca356 100644 --- a/src/memos/templates/mem_reader_prompts.py +++ b/src/memos/templates/mem_reader_prompts.py @@ -228,16 +228,27 @@ Your task is to process a document chunk and generate a single, structured JSON object. Please perform: -1. Identify key information that reflects factual content, insights, decisions, or implications from the documents — including any notable themes, conclusions, or data points. Allow a reader to fully understand the essence of the chunk without reading the original text. -2. Resolve all time, person, location, and event references clearly: +1. Organize the document chunk into useful retrieval memories using the smallest complete record as the unit. The memories should be minimally redundant and collectively cover the important source information. + - A memory should answer one natural question on its own. Keep attributes, qualifiers, values, and relationships that belong to the same record together. + - Split content when it contains distinct subjects, rules, conditions, decisions, events, or data records that are independently useful for retrieval. + - But do not split one fact, rule, record, or short list into fragments merely because it contains multiple fields or peer items. +2. Treat tables, lists, statistics, measurements, prices, dates, percentages, rankings, identifiers, and other structured or numerical content as high-priority information. + - Preserve every row and the relationship between its row label and column values, including units, currencies, signs, ranges, and original category or rank names. + - For a data table, normally create one memory per complete row and include the row label plus all related column values in that memory. Do not split one row into separate memories for individual cells or columns. + - Keep a simple list of peer items together when the items share one subject and have no item-specific conditions, values, or explanations. Split list entries only when they carry independently meaningful details. + - Never replace exact values or lower rows with vague summaries such as "and so on", "decreasing progressively", or "other entries are similar". + - Do not calculate, normalize, rename, reorder, or "correct" values and labels unless the source explicitly provides the correction. +3. Identify key information that reflects factual content, insights, decisions, or implications from the documents — including any notable themes, conclusions, or data points. Allow a reader to fully understand the essence of the chunk without reading the original text. +4. Resolve all time, person, location, and event references clearly: - Convert relative time expressions (e.g., “last year,” “next quarter”) into absolute dates if context allows. - Clearly distinguish between event time and document time. - If uncertainty exists, state it explicitly (e.g., “around 2024,” “exact date unclear”). + - Never invent a missing year or date. Use document-level context only when it unambiguously identifies the same event or subject. - Include specific locations if mentioned. - Resolve all pronouns, aliases, and ambiguous references into full names or identities. - Disambiguate entities with the same name if applicable. -3. Always write from a third-person perspective, referring to the subject or content clearly rather than using first-person ("I", "me", "my"). -4. Do not omit any information that is likely to be important or memorable from the document summaries. +5. Always write from a third-person perspective, referring to the subject or content clearly rather than using first-person ("I", "me", "my"). +6. Do not omit any information that is likely to be important or memorable from the document summaries. - Include all key facts, insights, emotional tones, and plans — even if they seem minor. - Prioritize completeness and fidelity over conciseness. - Do not generalize or skip details that could be contextually meaningful. @@ -249,7 +260,7 @@ { "key": , "memory_type": "LongTermMemory", - "value": , + "value": , "tags": } ... @@ -263,7 +274,7 @@ {custom_tags_prompt} -If given context, use it as a supplement to the document information extraction; if no context is given, directly process the document information. +Use the reference context only to disambiguate the current chunk, such as identifying the document, event, subject, or year. Do not create a memory solely from the reference context unless the same fact is supported by the current document chunk. If no context is given, directly process the document information. Reference context: {context} @@ -276,16 +287,27 @@ 您的任务是处理文档片段，并生成一个结构化的 JSON 列表对象。 请执行以下操作： -1. 识别反映文档中事实内容、见解、决策或含义的关键信息——包括任何显著的主题、结论或数据点，使读者无需阅读原文即可充分理解该片段的核心内容。 -2. 清晰解析所有时间、人物、地点和事件的指代： +1. 以“最小完整记录”为单位，将文档片段组织成便于检索的多条记忆；各条记忆应尽量不重复，并合起来覆盖原文中的重要信息。 + - 一条记忆应能独立回答一个自然问题；属于同一记录的属性、限定条件、数值和关系应保留在一起。 + - 当内容涉及不同主题、规则、条件、决策、事件或数据记录，并且分别具有独立检索价值时，需要进行拆分。 + - 但不要仅仅因为一个事实、规则、记录或简短列表包含多个字段或并列项，就把它机械拆成多个片段。 +2. 对表格、列表、统计数据、计量值、价格、日期、百分比、排名、编号以及其他结构化或数值信息给予额外关注。 + - 完整保留每一行，以及行名称与各列数值之间的对应关系，包括单位、币种、正负号、范围和原始类别或名次名称。 + - 对数据表格，通常每一行生成一条记忆，并在该条记忆中同时保留行名称及所有相关列值；不得把同一行拆成多条记忆，分别保存不同单元格或列。 + - 当普通并列列表中的各项属于同一主题，且没有各自独立的条件、数值或解释时，普通并列列表应按共同主题合并为一条记忆；只有列表项包含可独立检索的重要细节时才拆分。 + - 不得用模糊概述替代精确数值或后续行，例如“依次递减”“等等”“其他项目类似”。 + - 除非原文明确给出修正，否则不得自行计算、归一化、改名、重排或“纠正”数值和标签。 +3. 识别反映文档中事实内容、见解、决策或含义的关键信息——包括任何显著的主题、结论或数据点，使读者无需阅读原文即可充分理解该片段的核心内容。 +4. 清晰解析所有时间、人物、地点和事件的指代： - 如果上下文允许，将相对时间表达（如“去年”、“下一季度”）转换为绝对日期。 - 明确区分事件时间和文档时间。 - 如果存在不确定性，需明确说明（例如，“约2024年”，“具体日期不详”）。 + - 不得编造缺失的年份或日期。只有当文档级上下文能够无歧义地指向同一事件或主题时，才可据此补全年份。 - 若提及具体地点，请包含在内。 - 将所有代词、别名和模糊指代解析为全名或明确身份。 - 如有同名实体，需加以区分。 -3. 始终以第三人称视角撰写，清晰指代主题或内容，避免使用第一人称（“我”、“我们”、“我的”）。 -4. 不要遗漏文档摘要中可能重要或值得记忆的任何信息。 +5. 始终以第三人称视角撰写，清晰指代主题或内容，避免使用第一人称（“我”、“我们”、“我的”）。 +6. 不要遗漏文档摘要中可能重要或值得记忆的任何信息。 - 包括所有关键事实、见解、情感基调和计划——即使看似微小。 - 优先考虑完整性和保真度，而非简洁性。 - 不要泛化或跳过可能具有上下文意义的细节。 @@ -297,7 +319,7 @@ { "key": <字符串，`value` 字段的简洁标题>, "memory_type": "LongTermMemory", - "value": <一段清晰准确的段落，全面总结文档片段中的主要观点、论据和信息——若输入摘要为英文，则用英文；若为中文，则用中文>, + "value": <一条清晰、准确、自包含的最小完整记录，并保留紧密相关的字段、精确结构化数值及其对应关系——若输入摘要为英文，则用英文；若为中文，则用中文>, "tags": <相关主题关键词列表（例如，["截止日期", "团队", "计划"]）> } ... @@ -311,7 +333,7 @@ {custom_tags_prompt} -如果给定了上下文，就结合上下文信息作为文档信息提取的补充，如果没有给定上下文，请直接处理文档信息。 +参考上下文只用于消除当前片段中的歧义，例如确定文档、事件、主题或年份。除非当前文档片段本身也支持同一事实，否则不得仅依据参考上下文创建记忆。如果没有上下文，请直接处理文档信息。 参考的上下文： {context} diff --git a/tests/api/test_llm_provider_config.py b/tests/api/test_llm_provider_config.py index 61292c438..fbce5812b 100644 --- a/tests/api/test_llm_provider_config.py +++ b/tests/api/test_llm_provider_config.py @@ -32,6 +32,57 @@ def test_task_openai_model_uses_openai_provider_env(monkeypatch): assert config["config"]["api_base"] == "https://openai.example/v1" +def test_document_parser_model_uses_provider_env_and_dedicated_temperature(monkeypatch): + monkeypatch.setenv("DOCUMENT_PARSER_MODEL", "qwen3.6-flash") + monkeypatch.setenv("QWEN_API_KEY", "qwen-key") + monkeypatch.setenv("QWEN_API_BASE", "https://dashscope.example/v1") + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + monkeypatch.setenv("OPENAI_API_BASE", "https://openai.example/v1") + + config = APIConfig.get_document_parser_llm_config() + + assert config["backend"] == "qwen" + assert config["config"]["model_name_or_path"] == "qwen3.6-flash" + assert config["config"]["api_key"] == "qwen-key" + assert config["config"]["api_base"] == "https://dashscope.example/v1" + assert config["config"]["temperature"] == 0.8 + assert config["config"]["extra_body"] == {"enable_thinking": False} + + +def test_document_parser_model_falls_back_to_general_model(monkeypatch): + monkeypatch.delenv("DOCUMENT_PARSER_MODEL", raising=False) + monkeypatch.setenv("MEMREADER_GENERAL_MODEL", "gpt-4.1-mini") + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + monkeypatch.setenv("OPENAI_API_BASE", "https://openai.example/v1") + + config = APIConfig.get_document_parser_llm_config() + + assert config["backend"] == "openai" + assert config["config"]["model_name_or_path"] == "gpt-4.1-mini" + assert config["config"]["temperature"] == 0.8 + + +def test_document_parser_model_does_not_fall_back_to_main_memreader(monkeypatch): + monkeypatch.delenv("DOCUMENT_PARSER_MODEL", raising=False) + monkeypatch.delenv("MEMREADER_GENERAL_MODEL", raising=False) + monkeypatch.setenv("MEMRADER_MODEL", "qwen3-0.6B") + + assert APIConfig.get_document_parser_llm_config() is None + + +def test_product_config_wires_document_parser_model(monkeypatch): + monkeypatch.setenv("DOCUMENT_PARSER_MODEL", "gpt-4.1-mini") + monkeypatch.setenv("OPENAI_API_KEY", "openai-key") + monkeypatch.setenv("OPENAI_API_BASE", "https://openai.example/v1") + + reader_config = APIConfig.get_product_default_config()["mem_reader"]["config"] + + document_config = reader_config["document_parser_llm"] + assert document_config["backend"] == "openai" + assert document_config["config"]["model_name_or_path"] == "gpt-4.1-mini" + assert document_config["config"]["temperature"] == 0.8 + + def test_qwen_llm_only_uses_model_and_provider_endpoint_env(monkeypatch): monkeypatch.setenv("QWEN_MODEL", "qwen-flash") monkeypatch.setenv("QWEN_API_KEY", "qwen-key") diff --git a/tests/mem_reader/test_document_parser_llm.py b/tests/mem_reader/test_document_parser_llm.py new file mode 100644 index 000000000..1c40a321c --- /dev/null +++ b/tests/mem_reader/test_document_parser_llm.py @@ -0,0 +1,128 @@ +from unittest.mock import MagicMock, patch + +from memos.mem_reader.read_multi_modal.file_content_parser import FileContentParser +from memos.mem_reader.read_multi_modal.multi_modal_parser import MultiModalParser +from memos.templates.mem_reader_prompts import ( + SIMPLE_STRUCT_DOC_READER_PROMPT, + SIMPLE_STRUCT_DOC_READER_PROMPT_ZH, +) + + +def test_url_path_markdown_suffix_overrides_original_filename(): + parser = FileContentParser( + embedder=MagicMock(), + llm=MagicMock(), + direct_markdown_hostnames=[], + ) + response = MagicMock() + response.text = "# parsed markdown" + response.content = b"# parsed markdown" + + with ( + patch("requests.get", return_value=response), + patch( + "tempfile.NamedTemporaryFile", + side_effect=AssertionError("markdown URL should not be written to a temp file"), + ), + ): + text, temp_path, is_markdown = parser._handle_url( + "https://memos.example/api/download/document.md?signature=secret", + "original-document.docx", + ) + + assert text == "# parsed markdown" + assert temp_path is None + assert is_markdown is True + + +def test_multi_modal_parser_routes_files_to_dedicated_document_llm(): + main_llm = MagicMock(name="main_llm") + document_llm = MagicMock(name="document_llm") + + parser = MultiModalParser( + embedder=MagicMock(), + llm=main_llm, + document_parser_llm=document_llm, + ) + + assert parser.string_parser.llm is main_llm + assert parser.file_content_parser.llm is document_llm + + +def test_multi_modal_parser_does_not_fall_back_files_to_main_llm(): + main_llm = MagicMock(name="main_llm") + + parser = MultiModalParser( + embedder=MagicMock(), + llm=main_llm, + document_parser_llm=None, + ) + + assert parser.string_parser.llm is main_llm + assert parser.file_content_parser.llm is None + + +def test_document_prompts_use_complete_records_without_over_splitting(): + assert "smallest complete record" in SIMPLE_STRUCT_DOC_READER_PROMPT + assert "Preserve every row" in SIMPLE_STRUCT_DOC_READER_PROMPT + assert "Do not split one row into separate memories" in SIMPLE_STRUCT_DOC_READER_PROMPT + assert "Keep a simple list of peer items together" in SIMPLE_STRUCT_DOC_READER_PROMPT + assert "Never replace exact values" in SIMPLE_STRUCT_DOC_READER_PROMPT + + assert "最小完整记录" in SIMPLE_STRUCT_DOC_READER_PROMPT_ZH + assert "完整保留每一行" in SIMPLE_STRUCT_DOC_READER_PROMPT_ZH + assert "不得把同一行拆成多条记忆" in SIMPLE_STRUCT_DOC_READER_PROMPT_ZH + assert "普通并列列表应按共同主题合并" in SIMPLE_STRUCT_DOC_READER_PROMPT_ZH + assert "不得用模糊概述替代精确数值" in SIMPLE_STRUCT_DOC_READER_PROMPT_ZH + + +def test_markdown_h1_titles_are_passed_as_document_context(): + embedder = MagicMock() + embedder.embed.return_value = [[0.1]] + parser = FileContentParser( + embedder=embedder, + llm=MagicMock(), + parser=MagicMock(), + direct_markdown_hostnames=[], + ) + markdown = """# *2026 PPA 北京公开赛内部培训资料* + +## 非全局子标题 + +# 04 报名时间 + +即日起至 5 月 29 日 24:00 +""" + parser._handle_url = MagicMock(return_value=(markdown, None, True)) + parser._split_text = MagicMock(return_value=["# 04 报名时间\n即日起至 5 月 29 日 24:00"]) + parser._get_doc_llm_response = MagicMock( + return_value={ + "memory list": [ + { + "key": "报名截止时间", + "memory_type": "LongTermMemory", + "value": "报名截止时间为2026年5月29日24:00。", + "tags": ["报名", "时间"], + } + ], + "summary": "报名截止时间", + } + ) + + parser.parse_fine( + { + "type": "file", + "file": { + "file_data": "https://example.com/document.md", + "file_id": "file-1", + "filename": "training.docx", + }, + }, + {"user_id": "user-1", "session_id": "session-1"}, + ) + + context = parser._get_doc_llm_response.call_args.kwargs["message_text_context"] + assert "training.docx" in context + assert "*2026 PPA 北京公开赛内部培训资料*" in context + assert "04 报名时间" in context + assert "非全局子标题" not in context