feat: update file content memory extract

Author: fridayL

src/memos/mem_reader/read_multi_modal/file_content_parser.py

@@ -1,10 +1,14 @@
 """Parser for file content parts (RawMessageList)."""
 
+import concurrent.futures
 import os
 import tempfile
 
 from typing import Any
 
+from tqdm import tqdm
+
+from memos.context.context import ContextThreadPoolExecutor
 from memos.embedders.base import BaseEmbedder
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
@@ -13,19 +17,72 @@
     TextualMemoryItem,
     TreeNodeTextualMemoryMetadata,
 )
+from memos.templates.mem_reader_prompts import (
+    CUSTOM_TAGS_INSTRUCTION,
+    CUSTOM_TAGS_INSTRUCTION_ZH,
+    SIMPLE_STRUCT_DOC_READER_PROMPT,
+    SIMPLE_STRUCT_DOC_READER_PROMPT_ZH,
+)
 from memos.types.openai_chat_completion_types import File
 
 from .base import BaseMessageParser, _derive_key
-from .utils import get_parser, get_text_splitter
+from .utils import (
+    detect_lang,
+    get_parser,
+    get_text_splitter,
+    parse_json_result,
+)
 
 
 logger = get_logger(__name__)
 
+# Prompt dictionary for doc processing (shared by simple_struct and file_content_parser)
+DOC_PROMPT_DICT = {
+    "doc": {"en": SIMPLE_STRUCT_DOC_READER_PROMPT, "zh": SIMPLE_STRUCT_DOC_READER_PROMPT_ZH},
+    "custom_tags": {"en": CUSTOM_TAGS_INSTRUCTION, "zh": CUSTOM_TAGS_INSTRUCTION_ZH},
+}
+
 
 class FileContentParser(BaseMessageParser):
     """Parser for file content parts."""
 
-    def _handle_url(self, url_str: str, filename: str) -> tuple[str, str | None]:
+    def _get_doc_llm_response(self, chunk_text: str, custom_tags: list[str] | None = None) -> dict:
+        """
+        Call LLM to extract memory from document chunk.
+        Uses doc prompts from DOC_PROMPT_DICT.
+
+        Args:
+            chunk_text: Text chunk to extract memory from
+            custom_tags: Optional list of custom tags for LLM extraction
+
+        Returns:
+            Parsed JSON response from LLM or empty dict if failed
+        """
+        if not self.llm:
+            logger.warning("[FileContentParser] LLM not available for fine mode")
+            return {}
+
+        lang = detect_lang(chunk_text)
+        template = DOC_PROMPT_DICT["doc"][lang]
+        prompt = template.replace("{chunk_text}", chunk_text)
+
+        custom_tags_prompt = (
+            DOC_PROMPT_DICT["custom_tags"][lang].replace("{custom_tags}", str(custom_tags))
+            if custom_tags
+            else ""
+        )
+        prompt = prompt.replace("{custom_tags_prompt}", custom_tags_prompt)
+
+        messages = [{"role": "user", "content": prompt}]
+        try:
+            response_text = self.llm.generate(messages)
+            response_json = parse_json_result(response_text)
+        except Exception as e:
+            logger.error(f"[FileContentParser] LLM generation error: {e}")
+            response_json = {}
+        return response_json
+
+    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
@@ -42,14 +99,14 @@ def _handle_url(self, url_str: str, filename: str) -> tuple[str, str | None]:
                 filename = os.path.basename(parsed_url.path) or "downloaded_file"
 
             if hostname in self.direct_markdown_hostnames:
-                return response.text, None
+                return response.text, None, True
 
             file_ext = os.path.splitext(filename)[1].lower()
             if file_ext in [".md", ".markdown", ".txt"]:
-                return response.text, None
+                return response.text, None, True
             with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=file_ext) as temp_file:
                 temp_file.write(response.content)
-            return "", temp_file.name
+            return "", temp_file.name, False
         except Exception as e:
             logger.error(f"[FileContentParser] URL processing error: {e}")
             return f"[File URL download failed: {url_str}]", None
@@ -108,7 +165,7 @@ def __init__(
             else:
                 self.direct_markdown_hostnames = []
 
-    def _split_text(self, text: str) -> list[str]:
+    def _split_text(self, text: str, is_markdown: bool = False) -> list[str]:
         """
         Split text into chunks using text splitter from utils.
 
@@ -288,6 +345,8 @@ def parse_fast(
 
         # Extract info fields
         info_ = info.copy()
+        if file_id:
+            info_.update({"file_id": file_id})
         user_id = info_.pop("user_id", "")
         session_id = info_.pop("session_id", "")
 
@@ -358,10 +417,19 @@ def parse_fine(
         """
         Parse file content part in fine mode.
         Fine mode downloads and parses file content, especially for URLs.
+        Then uses LLM to extract structured memories from each chunk.
+
         Handles various file parameter scenarios:
         - file_data: URL (http://, https://, or @http://), base64 encoded data, or plain text content
         - file_id: ID of an uploaded file
         - filename: name of the file
+
+        Args:
+            message: File content part to parse
+            info: Dictionary containing user_id and session_id
+            **kwargs: Additional parameters including:
+                - custom_tags: Optional list of custom tags for LLM extraction
+                - context_items: Optional list of TextualMemoryItem for context
         """
         if not isinstance(message, dict):
             logger.warning(f"[FileContentParser] Expected dict, got {type(message)}")
@@ -378,6 +446,9 @@ def parse_fine(
         file_id = file_info.get("file_id", "")
         filename = file_info.get("filename", "")
 
+        # Extract custom_tags from kwargs (for LLM extraction)
+        custom_tags = kwargs.get("custom_tags")
+
         # Use parser from utils
         parser = self.parser or get_parser()
         if not parser:
@@ -386,6 +457,7 @@ def parse_fine(
 
         parsed_text = ""
         temp_file_path = None
+        is_markdown = False
 
         try:
             # Priority 1: If file_data is provided, process it
@@ -394,7 +466,9 @@ def parse_fine(
                     url_str = file_data[1:] if file_data.startswith("@") else file_data
 
                     if url_str.startswith(("http://", "https://")):
-                        parsed_text, temp_file_path = self._handle_url(url_str, filename)
+                        parsed_text, temp_file_path, is_markdown = self._handle_url(
+                            url_str, filename
+                        )
                         if temp_file_path:
                             try:
                                 # Use parser from utils
@@ -457,28 +531,32 @@ def parse_fine(
         memory_type = "LongTermMemory"
 
         # Split parsed text into chunks
-        content_chunks = self._split_text(parsed_text)
-
-        # Create memory items for each chunk
-        memory_items = []
-        for chunk_idx, chunk_text in enumerate(content_chunks):
-            if not chunk_text.strip():
-                continue
-
-            memory_item = TextualMemoryItem(
-                memory=chunk_text,
+        content_chunks = self._split_text(parsed_text, is_markdown)
+
+        # Filter out empty chunks and create indexed list
+        valid_chunks = [
+            (idx, chunk_text) for idx, chunk_text in enumerate(content_chunks) if chunk_text.strip()
+        ]
+        total_chunks = len(content_chunks)
+
+        # Helper function to create memory item (similar to SimpleStructMemReader._make_memory_item)
+        def _make_memory_item(
+            value: str,
+            mem_type: str = memory_type,
+            tags: list[str] | None = None,
+            key: str | None = None,
+        ) -> TextualMemoryItem:
+            """Construct memory item with common fields."""
+            return TextualMemoryItem(
+                memory=value,
                 metadata=TreeNodeTextualMemoryMetadata(
                     user_id=user_id,
                     session_id=session_id,
-                    memory_type=memory_type,
+                    memory_type=mem_type,
                     status="activated",
-                    tags=[
-                        "mode:fine",
-                        "multimodal:file",
-                        f"chunk:{chunk_idx + 1}/{len(content_chunks)}",
-                    ],
-                    key=_derive_key(chunk_text),
-                    embedding=self.embedder.embed([chunk_text])[0],
+                    tags=tags or [],
+                    key=key if key is not None else _derive_key(value),
+                    embedding=self.embedder.embed([value])[0],
                     usage=[],
                     sources=[source],
                     background="",
@@ -487,28 +565,102 @@ def parse_fine(
                     info=info_,
                 ),
             )
-            memory_items.append(memory_item)
 
-        # If no chunks were created, create a placeholder
-        if not memory_items:
-            memory_item = TextualMemoryItem(
-                memory=parsed_text,
-                metadata=TreeNodeTextualMemoryMetadata(
-                    user_id=user_id,
-                    session_id=session_id,
-                    memory_type=memory_type,
-                    status="activated",
-                    tags=["mode:fine", "multimodal:file"],
-                    key=_derive_key(parsed_text),
-                    embedding=self.embedder.embed([parsed_text])[0],
-                    usage=[],
-                    sources=[source],
-                    background="",
-                    confidence=0.99,
-                    type="fact",
-                    info=info_,
-                ),
+        # Helper function to create fallback item for a chunk
+        def _make_fallback(
+            chunk_idx: int, chunk_text: str, reason: str = "raw"
+        ) -> TextualMemoryItem:
+            """Create fallback memory item with raw chunk text."""
+            return _make_memory_item(
+                value=chunk_text,
+                tags=[
+                    "mode:fine",
+                    "multimodal:file",
+                    f"fallback:{reason}",
+                    f"chunk:{chunk_idx + 1}/{total_chunks}",
+                ],
             )
-            memory_items.append(memory_item)
 
-        return memory_items
+        # Handle empty chunks case
+        if not valid_chunks:
+            return [
+                _make_memory_item(
+                    value=parsed_text or "[File: empty content]",
+                    tags=["mode:fine", "multimodal:file"],
+                )
+            ]
+
+        # If no LLM available, create memory items directly from chunks
+        if not self.llm:
+            return [_make_fallback(idx, text, "no_llm") for idx, text in valid_chunks]
+
+        # Process single chunk with LLM extraction (worker function)
+        def _process_chunk(chunk_idx: int, chunk_text: str) -> TextualMemoryItem:
+            """Process chunk with LLM, fallback to raw on failure."""
+            try:
+                response_json = self._get_doc_llm_response(chunk_text, custom_tags)
+                if response_json:
+                    value = response_json.get("value", "").strip()
+                    if value:
+                        tags = response_json.get("tags", [])
+                        tags = tags if isinstance(tags, list) else []
+                        tags.extend(["mode:fine", "multimodal:file"])
+
+                        llm_mem_type = response_json.get("memory_type", memory_type)
+                        if llm_mem_type not in ["LongTermMemory", "UserMemory"]:
+                            llm_mem_type = memory_type
+
+                        return _make_memory_item(
+                            value=value,
+                            mem_type=llm_mem_type,
+                            tags=tags,
+                            key=response_json.get("key"),
+                        )
+            except Exception as e:
+                logger.error(f"[FileContentParser] LLM error for chunk {chunk_idx}: {e}")
+
+            # Fallback to raw chunk
+            logger.warning(f"[FileContentParser] Fallback to raw for chunk {chunk_idx}")
+            return _make_fallback(chunk_idx, chunk_text)
+
+        # Process chunks concurrently with progress bar
+        memory_items = []
+        chunk_map = dict(valid_chunks)
+        total_chunks = len(valid_chunks)
+
+        logger.info(f"[FileContentParser] Processing {total_chunks} chunks with LLM...")
+
+        with ContextThreadPoolExecutor(max_workers=20) as executor:
+            futures = {
+                executor.submit(_process_chunk, idx, text): idx for idx, text in valid_chunks
+            }
+
+            # Use tqdm for progress bar (similar to simple_struct.py _process_doc_data)
+            for future in tqdm(
+                concurrent.futures.as_completed(futures),
+                total=total_chunks,
+                desc="[FileContentParser] Processing chunks",
+            ):
+                chunk_idx = futures[future]
+                try:
+                    node = future.result()
+                    if node:
+                        memory_items.append(node)
+                except Exception as e:
+                    tqdm.write(f"[ERROR] Chunk {chunk_idx} failed: {e}")
+                    logger.error(f"[FileContentParser] Future failed for chunk {chunk_idx}: {e}")
+                    # Create fallback for failed future
+                    if chunk_idx in chunk_map:
+                        memory_items.append(
+                            _make_fallback(chunk_idx, chunk_map[chunk_idx], "error")
+                        )
+
+        logger.info(
+            f"[FileContentParser] Completed processing {len(memory_items)}/{total_chunks} chunks"
+        )
+
+        return memory_items or [
+            _make_memory_item(
+                value=parsed_text or "[File: empty content]", tags=["mode:fine", "multimodal:file"]
+            )
+        ]