MemoryOS 2.0.3__py3-none-any.whl

memos/mem_reader/read_multi_modal/image_parser.py

@@ -0,0 +1,359 @@
+"""Parser for image_url content parts."""
+
+import json
+import re
+
+from typing import Any
+
+from memos.embedders.base import BaseEmbedder
+from memos.llms.base import BaseLLM
+from memos.log import get_logger
+from memos.memories.textual.item import (
+    SourceMessage,
+    TextualMemoryItem,
+    TreeNodeTextualMemoryMetadata,
+)
+from memos.templates.mem_reader_prompts import IMAGE_ANALYSIS_PROMPT_EN, IMAGE_ANALYSIS_PROMPT_ZH
+from memos.types.openai_chat_completion_types import ChatCompletionContentPartImageParam
+
+from .base import BaseMessageParser, _derive_key
+from .utils import detect_lang
+
+
+logger = get_logger(__name__)
+
+
+class ImageParser(BaseMessageParser):
+    """Parser for image_url content parts."""
+
+    def __init__(self, embedder: BaseEmbedder, llm: BaseLLM | None = None):
+        """
+        Initialize ImageParser.
+
+        Args:
+            embedder: Embedder for generating embeddings
+            llm: Optional LLM for fine mode processing
+        """
+        super().__init__(embedder, llm)
+
+    def create_source(
+        self,
+        message: ChatCompletionContentPartImageParam,
+        info: dict[str, Any],
+    ) -> SourceMessage:
+        """Create SourceMessage from image_url content part."""
+        if isinstance(message, dict):
+            image_url = message.get("image_url", {})
+            if isinstance(image_url, dict):
+                url = image_url.get("url", "")
+                detail = image_url.get("detail", "auto")
+            else:
+                url = str(image_url)
+                detail = "auto"
+            return SourceMessage(
+                type="image",
+                content=url,
+                url=url,
+                detail=detail,
+            )
+        return SourceMessage(type="image", content=str(message))
+
+    def rebuild_from_source(
+        self,
+        source: SourceMessage,
+    ) -> ChatCompletionContentPartImageParam:
+        """Rebuild image_url content part from SourceMessage."""
+        # Rebuild from source fields
+        url = (
+            getattr(source, "url", "")
+            or getattr(source, "image_path", "")
+            or (source.content or "").replace("[image_url]: ", "")
+        )
+        detail = getattr(source, "detail", "auto")
+        return {
+            "type": "image_url",
+            "image_url": {
+                "url": url,
+                "detail": detail,
+            },
+        }
+
+    def parse_fast(
+        self,
+        message: ChatCompletionContentPartImageParam,
+        info: dict[str, Any],
+        **kwargs,
+    ) -> list[TextualMemoryItem]:
+        """Parse image_url in fast mode - returns empty list as images need fine mode processing."""
+        # In fast mode, images are not processed (they need vision models)
+        # They will be processed in fine mode via process_transfer
+        return []
+
+    def parse_fine(
+        self,
+        message: ChatCompletionContentPartImageParam,
+        info: dict[str, Any],
+        **kwargs,
+    ) -> list[TextualMemoryItem]:
+        """
+        Parse image_url in fine mode using vision models to extract information from images.
+
+        Args:
+            message: Image message to parse
+            info: Dictionary containing user_id and session_id
+            **kwargs: Additional parameters (e.g., context_items, custom_tags)
+
+        Returns:
+            List of TextualMemoryItem objects extracted from the image
+        """
+        if not self.llm:
+            logger.warning("[ImageParser] LLM not available for fine mode processing")
+            return []
+
+        # Extract image information
+        if not isinstance(message, dict):
+            logger.warning(f"[ImageParser] Expected dict, got {type(message)}")
+            return []
+
+        image_url = message.get("image_url", {})
+        if isinstance(image_url, dict):
+            url = image_url.get("url", "")
+            detail = image_url.get("detail", "auto")
+        else:
+            url = str(image_url)
+            detail = "auto"
+
+        if not url:
+            logger.warning("[ImageParser] No image URL found in message")
+            return []
+
+        # Create source for this image
+        source = self.create_source(message, info)
+
+        # Get context items if available
+        context_items = kwargs.get("context_items")
+
+        # Determine language: prioritize lang from source (passed via kwargs),
+        # fallback to detecting from context_items if lang not provided
+        lang = kwargs.get("lang")
+        if lang is None and context_items:
+            for item in context_items:
+                if hasattr(item, "memory") and item.memory:
+                    lang = detect_lang(item.memory)
+                    break
+        if not lang:
+            lang = "en"
+        if not hasattr(source, "lang") or source.lang is None:
+            source.lang = lang
+
+        # Select prompt based on language
+        image_analysis_prompt = (
+            IMAGE_ANALYSIS_PROMPT_ZH if lang == "zh" else IMAGE_ANALYSIS_PROMPT_EN
+        )
+
+        # Add context if available
+        context_text = ""
+        if context_items:
+            for item in context_items:
+                if hasattr(item, "memory") and item.memory:
+                    context_text += f"{item.memory}\n"
+        context_text = context_text.strip()
+
+        # Inject context into prompt when possible
+        image_analysis_prompt = image_analysis_prompt.replace("{context}", context_text)
+
+        # Build messages with image content
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": image_analysis_prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": url,
+                            "detail": detail,
+                        },
+                    },
+                ],
+            }
+        ]
+
+        try:
+            # Call LLM with vision model
+            response_text = self.llm.generate(messages)
+            if not response_text:
+                logger.warning("[ImageParser] Empty response from LLM")
+                return []
+
+            # Parse JSON response
+            response_json = self._parse_json_result(response_text)
+            if not response_json:
+                logger.warning(f"[ImageParser] Fail to parse response from LLM: {response_text}")
+                return []
+
+            # Extract memory items from response
+            memory_items = []
+            memory_list = response_json.get("memory list", [])
+
+            if not memory_list:
+                logger.warning("[ImageParser] No memory items extracted from image")
+                # Fallback: create a simple memory item with the summary
+                summary = response_json.get(
+                    "summary", "Image analyzed but no specific memories extracted."
+                )
+                if summary:
+                    memory_items.append(
+                        self._create_memory_item(
+                            value=summary,
+                            info=info,
+                            memory_type="LongTermMemory",
+                            tags=["image", "visual"],
+                            key=_derive_key(summary),
+                            sources=[source],
+                            background=summary,
+                        )
+                    )
+                return memory_items
+
+            # Create memory items from parsed response
+            for mem_data in memory_list:
+                try:
+                    # Normalize memory_type
+                    memory_type = (
+                        mem_data.get("memory_type", "LongTermMemory")
+                        .replace("长期记忆", "LongTermMemory")
+                        .replace("用户记忆", "UserMemory")
+                    )
+                    if memory_type not in ["LongTermMemory", "UserMemory"]:
+                        memory_type = "LongTermMemory"
+
+                    value = mem_data.get("value", "").strip()
+                    if not value:
+                        continue
+
+                    tags = mem_data.get("tags", [])
+                    if not isinstance(tags, list):
+                        tags = []
+                    # Add image-related tags
+                    if "image" not in [t.lower() for t in tags]:
+                        tags.append("image")
+                    if "visual" not in [t.lower() for t in tags]:
+                        tags.append("visual")
+
+                    key = mem_data.get("key", "")
+                    background = response_json.get("summary", "")
+
+                    memory_item = self._create_memory_item(
+                        value=value,
+                        info=info,
+                        memory_type=memory_type,
+                        tags=tags,
+                        key=key if key else _derive_key(value),
+                        sources=[source],
+                        background=background,
+                    )
+                    memory_items.append(memory_item)
+                except Exception as e:
+                    logger.error(f"[ImageParser] Error creating memory item: {e}")
+                    continue
+
+            return memory_items
+
+        except Exception as e:
+            logger.error(f"[ImageParser] Error processing image in fine mode: {e}")
+            # Fallback: create a simple memory item
+            fallback_value = f"Image analyzed: {url}"
+            return [
+                self._create_memory_item(
+                    value=fallback_value,
+                    info=info,
+                    memory_type="LongTermMemory",
+                    tags=["image", "visual"],
+                    key=_derive_key(fallback_value),
+                    sources=[source],
+                    background="Image processing encountered an error.",
+                )
+            ]
+
+    def _parse_json_result(self, response_text: str) -> dict:
+        """
+        Parse JSON result from LLM response.
+        Similar to SimpleStructMemReader.parse_json_result.
+        """
+        s = (response_text or "").strip()
+
+        # Try to extract JSON from code blocks
+        m = re.search(r"```(?:json)?\s*([\s\S]*?)```", s, flags=re.I)
+        s = (m.group(1) if m else s.replace("```", "")).strip()
+
+        # Find first {
+        i = s.find("{")
+        if i == -1:
+            return {}
+        s = s[i:].strip()
+
+        try:
+            return json.loads(s)
+        except json.JSONDecodeError:
+            pass
+
+        # Try to find the last } or ]
+        j = max(s.rfind("}"), s.rfind("]"))
+        if j != -1:
+            try:
+                return json.loads(s[: j + 1])
+            except json.JSONDecodeError:
+                pass
+
+        # Try to close brackets
+        def _cheap_close(t: str) -> str:
+            t += "}" * max(0, t.count("{") - t.count("}"))
+            t += "]" * max(0, t.count("[") - t.count("]"))
+            return t
+
+        t = _cheap_close(s)
+        try:
+            return json.loads(t)
+        except json.JSONDecodeError as e:
+            if "Invalid \\escape" in str(e):
+                s = s.replace("\\", "\\\\")
+                try:
+                    return json.loads(s)
+                except json.JSONDecodeError:
+                    pass
+            logger.warning(f"[ImageParser] Failed to parse JSON: {e}\nResponse: {response_text}")
+
+    def _create_memory_item(
+        self,
+        value: str,
+        info: dict[str, Any],
+        memory_type: str,
+        tags: list[str],
+        key: str,
+        sources: list[SourceMessage],
+        background: str = "",
+    ) -> TextualMemoryItem:
+        """Create a TextualMemoryItem with the given parameters."""
+        info_ = info.copy()
+        user_id = info_.pop("user_id", "")
+        session_id = info_.pop("session_id", "")
+
+        return TextualMemoryItem(
+            memory=value,
+            metadata=TreeNodeTextualMemoryMetadata(
+                user_id=user_id,
+                session_id=session_id,
+                memory_type=memory_type,
+                status="activated",
+                tags=tags,
+                key=key,
+                embedding=self.embedder.embed([value])[0],
+                usage=[],
+                sources=sources,
+                background=background,
+                confidence=0.99,
+                type="fact",
+                info=info_,
+            ),
+        )

memos/mem_reader/read_multi_modal/multi_modal_parser.py

@@ -0,0 +1,252 @@
+"""Unified multimodal parser for different message types.
+
+This module provides a unified interface to parse different message types
+in both fast and fine modes.
+"""
+
+from typing import Any
+
+from memos.embedders.base import BaseEmbedder
+from memos.llms.base import BaseLLM
+from memos.log import get_logger
+from memos.memories.textual.item import SourceMessage, TextualMemoryItem
+from memos.types import MessagesType
+
+from .assistant_parser import AssistantParser
+from .base import BaseMessageParser
+from .file_content_parser import FileContentParser
+from .image_parser import ImageParser
+from .string_parser import StringParser
+from .system_parser import SystemParser
+from .text_content_parser import TextContentParser
+from .tool_parser import ToolParser
+from .user_parser import UserParser
+from .utils import extract_role
+
+
+logger = get_logger(__name__)
+
+
+class MultiModalParser:
+    """Unified parser for different message types."""
+
+    def __init__(
+        self,
+        embedder: BaseEmbedder,
+        llm: BaseLLM | None = None,
+        parser: Any | None = None,
+        direct_markdown_hostnames: list[str] | None = None,
+    ):
+        """
+        Initialize MultiModalParser.
+
+        Args:
+            embedder: Embedder for generating embeddings
+            llm: Optional LLM for fine mode processing
+            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
+                environment variable (comma-separated). Default: ["139.196.232.20"]
+        """
+        self.embedder = embedder
+        self.llm = llm
+        self.parser = parser
+
+        # Initialize parsers for different message types
+        self.string_parser = StringParser(embedder, llm)
+        self.system_parser = SystemParser(embedder, llm)
+        self.user_parser = UserParser(embedder, llm)
+        self.assistant_parser = AssistantParser(embedder, llm)
+        self.tool_parser = ToolParser(embedder, llm)
+        self.text_content_parser = TextContentParser(embedder, llm)
+        self.file_content_parser = FileContentParser(
+            embedder, llm, parser, direct_markdown_hostnames=direct_markdown_hostnames
+        )
+        self.image_parser = ImageParser(embedder, llm)
+        self.audio_parser = None  # future
+
+        self.role_parsers = {
+            "system": SystemParser(embedder, llm),
+            "user": UserParser(embedder, llm),
+            "assistant": AssistantParser(embedder, llm),
+            "tool": ToolParser(embedder, llm),
+        }
+
+        self.type_parsers = {
+            "text": self.text_content_parser,
+            "file": self.file_content_parser,
+            "image": self.image_parser,
+            "image_url": self.image_parser,  # Support both "image" and "image_url"
+            "audio": self.audio_parser,
+            # Custom tool formats
+            "tool_description": self.tool_parser,
+            "tool_input": self.tool_parser,
+            "tool_output": self.tool_parser,
+        }
+
+    def _get_parser(self, message: Any) -> BaseMessageParser | None:
+        """
+        Get appropriate parser for the message type.
+
+        Args:
+            message: Message to parse
+
+        Returns:
+            Appropriate parser or None
+        """
+        # Handle string messages
+        if isinstance(message, str):
+            return self.string_parser
+
+        # Handle dict messages
+        if not isinstance(message, dict):
+            logger.warning(f"[MultiModalParser] Unknown message type: {type(message)}")
+            return None
+
+        # Check if it's a RawMessageList item (text or file)
+        if "type" in message:
+            msg_type = message.get("type")
+            parser = self.type_parsers.get(msg_type)
+            if parser:
+                return parser
+
+        # Check if it's a MessageList item (system, user, assistant, tool)
+        role = extract_role(message)
+        if role:
+            parser = self.role_parsers.get(role)
+            if parser:
+                return parser
+
+        logger.warning(f"[MultiModalParser] Could not determine parser for message: {message}")
+        return None
+
+    def parse(
+        self,
+        message: MessagesType,
+        info: dict[str, Any],
+        mode: str = "fast",
+        **kwargs,
+    ) -> list[TextualMemoryItem]:
+        """
+        Parse a single message in the specified mode.
+
+        Args:
+            message: Message to parse (can be str, MessageList item, or RawMessageList item)
+            info: Dictionary containing user_id and session_id
+            mode: "fast" or "fine"
+            **kwargs: Additional parameters
+
+        Returns:
+            List of TextualMemoryItem objects
+        """
+        # Handle list of messages (MessageList or RawMessageList)
+        if isinstance(message, list):
+            return [item for msg in message for item in self.parse(msg, info, mode, **kwargs)]
+
+        # Get appropriate parser
+        parser = self._get_parser(message)
+        if not parser:
+            logger.warning(f"[MultiModalParser] No parser found for message: {message}")
+            return []
+
+        # Parse using the appropriate parser
+        try:
+            return parser.parse(message, info, mode=mode, **kwargs)
+        except Exception as e:
+            logger.error(f"[MultiModalParser] Error parsing message: {e}")
+            return []
+
+    def parse_batch(
+        self,
+        messages: list[MessagesType],
+        info: dict[str, Any],
+        mode: str = "fast",
+        **kwargs,
+    ) -> list[list[TextualMemoryItem]]:
+        """
+        Parse a batch of messages.
+
+        Args:
+            messages: List of messages to parse
+            info: Dictionary containing user_id and session_id
+            mode: "fast" or "fine"
+            **kwargs: Additional parameters
+
+        Returns:
+            List of lists of TextualMemoryItem objects (one list per message)
+        """
+        results = []
+        for message in messages:
+            items = self.parse(message, info, mode, **kwargs)
+            results.append(items)
+        return results
+
+    def process_transfer(
+        self,
+        source: SourceMessage,
+        context_items: list[TextualMemoryItem] | None = None,
+        **kwargs,
+    ) -> list[TextualMemoryItem]:
+        """
+        Process transfer from SourceMessage to fine memory items.
+
+        This method:
+        1. Determines which parser to use based on source type
+        2. Rebuilds message from source using parser's rebuild_from_source
+        3. Calls parse_fine on the appropriate parser
+
+        Args:
+            source: SourceMessage to process
+            context_items: Optional list of TextualMemoryItem for context
+            **kwargs: Additional parameters (e.g., info dict with user_id, session_id, custom_tags)
+
+        Returns:
+            List of TextualMemoryItem objects from fine mode parsing
+        """
+        if not self.llm:
+            logger.warning("[MultiModalParser] LLM not available for process_transfer")
+            return []
+
+        # Extract info from context_items if available
+        info = kwargs.get("info", {})
+        if context_items and len(context_items) > 0:
+            first_item = context_items[0]
+            if not info:
+                info = {
+                    "user_id": first_item.metadata.user_id,
+                    "session_id": first_item.metadata.session_id,
+                }
+
+        # Try to determine parser from source.type
+        parser = None
+        if source.type == "file":
+            parser = self.file_content_parser
+        elif source.type == "text":
+            parser = self.text_content_parser
+        elif source.type in ["image", "image_url"]:
+            parser = self.image_parser
+        elif source.role:
+            # Chat message, use role parser
+            parser = self.role_parsers.get(source.role)
+
+        if not parser:
+            logger.warning(f"[MultiModalParser] Could not determine parser for source: {source}")
+            return []
+
+        # Rebuild message from source using parser's method
+        try:
+            message = parser.rebuild_from_source(source)
+        except Exception as e:
+            logger.error(f"[MultiModalParser] Error rebuilding message from source: {e}")
+            return []
+
+        # Parse in fine mode (pass context_items and custom_tags to parse_fine)
+        try:
+            custom_tags = kwargs.pop("custom_tags", None)
+            info = kwargs.pop("info", None)
+            return parser.parse_fine(
+                message, info, context_items=context_items, custom_tags=custom_tags, **kwargs
+            )
+        except Exception as e:
+            logger.error(f"[MultiModalParser] Error parsing in fine mode: {e}")
+            return []