modify multi-modal code

Author: yuan.wang

src/memos/mem_reader/read_multi_modal/system_parser.py

@@ -37,17 +37,21 @@ def __init__(self, embedder: BaseEmbedder, llm: BaseLLM | None = None):
 
     def create_source(
         self,
-        message: str,
+        message: ChatCompletionSystemMessageParam,
         info: dict[str, Any],
     ) -> SourceMessage:
         """Create SourceMessage from system message."""
+        content = message["content"]
+        if isinstance(content, dict):
+            content = content["text"]
+
         content_wo_tool_schema = re.sub(
             r"<tool_schema>(.*?)</tool_schema>",
             r"<tool_schema>omitted</tool_schema>",
-            message,
+            content,
             flags=re.DOTALL,
         )
-        tool_schema_match = re.search(r"<tool_schema>(.*?)</tool_schema>", message, re.DOTALL)
+        tool_schema_match = re.search(r"<tool_schema>(.*?)</tool_schema>", content, re.DOTALL)
         tool_schema_content = tool_schema_match.group(1) if tool_schema_match else ""
 
         return SourceMessage(
@@ -90,7 +94,7 @@ def parse_fast(
             flags=re.DOTALL,
         )
 
-        source = self.create_source(content, info)
+        source = self.create_source(message, info)
         return [
             TextualMemoryItem(
                 memory=content_wo_tool_schema,
@@ -125,9 +129,10 @@ def parse_fine(
         return [
             TextualMemoryItem(
                 id=str(uuid.uuid4()),
-                memory=json.dumps(tool_schema),
+                memory=json.dumps(schema),
                 metadata=TreeNodeTextualMemoryMetadata(
                     memory_type="ToolSchemaMemory",
                 ),
             )
+            for schema in tool_schema
         ]

src/memos/mem_reader/read_multi_modal/tool_parser.py

@@ -230,26 +230,28 @@ def parse_fast(
         content = message.get("content", "")
         chat_time = message.get("chat_time", None)
 
-        if role != "user":
-            logger.warning(f"[ToolParser] Expected role is `user`, got {role}")
+        if role != "tool":
+            logger.warning(f"[ToolParser] Expected role is `tool`, got {role}")
             return []
         parts = [f"{role}: "]
         if chat_time:
             parts.append(f"[{chat_time}]: ")
         prefix = "".join(parts)
-        content = json.dumps(content) if isinstance(content, list) else content
+        content = json.dumps(content) if isinstance(content, list | dict) else content
         line = f"{prefix}{content}\n"
         if not line:
             return []
-        memory_type = (
-            "LongTermMemory"  # only choce long term memory for tool messages as a placeholder
-        )
 
         sources = self.create_source(message, info)
         return [
             TextualMemoryItem(
                 memory=line,
-                metadata=TreeNodeTextualMemoryMetadata(memory_type=memory_type, sources=sources),
+                metadata=TreeNodeTextualMemoryMetadata(
+                    memory_type="LongTermMemory",  # only choce long term memory for tool messages as a placeholder
+                    status="activated",
+                    tags=["mode:fast"],
+                    sources=sources,
+                ),
             )
         ]
 

src/memos/memories/textual/tree_text_memory/organize/manager.py

@@ -186,7 +186,12 @@ def _process_memory(self, memory: TextualMemoryItem, user_name: str | None = Non
             )
             futures.append(("working", f_working))
 
-            if memory.metadata.memory_type in ("LongTermMemory", "UserMemory"):
+            if memory.metadata.memory_type in (
+                "LongTermMemory",
+                "UserMemory",
+                "ToolSchemaMemory",
+                "ToolTrajectoryMemory",
+            ):
                 f_graph = ex.submit(
                     self._add_to_graph_memory,
                     memory=memory,

src/memos/types/openai_chat_completion_types/chat_completion_assistant_message_param.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
 from typing import Literal, TypeAlias
 
 from typing_extensions import Required, TypedDict
@@ -35,7 +34,7 @@ class ChatCompletionAssistantMessageParam(TypedDict, total=False):
     [Learn more](https://platform.openai.com/docs/guides/audio).
     """
 
-    content: str | Iterable[ContentArrayOfContentPart] | None
+    content: str | list[ContentArrayOfContentPart] | ContentArrayOfContentPart | None
     """The contents of the assistant message.
 
     Required unless `tool_calls` or `function_call` is specified.
@@ -44,7 +43,9 @@ class ChatCompletionAssistantMessageParam(TypedDict, total=False):
     refusal: str | None
     """The refusal message by the assistant."""
 
-    tool_calls: Iterable[ChatCompletionMessageToolCallUnionParam]
+    tool_calls: (
+        list[ChatCompletionMessageToolCallUnionParam] | ChatCompletionMessageToolCallUnionParam
+    )
     """The tool calls generated by the model, such as function calls."""
 
     chat_time: str | None

src/memos/types/openai_chat_completion_types/chat_completion_system_message_param.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
 from typing import Literal
 
 from typing_extensions import Required, TypedDict
@@ -14,7 +13,9 @@
 
 
 class ChatCompletionSystemMessageParam(TypedDict, total=False):
-    content: Required[str | Iterable[ChatCompletionContentPartTextParam]]
+    content: Required[
+        str | list[ChatCompletionContentPartTextParam] | ChatCompletionContentPartTextParam
+    ]
     """The contents of the system message."""
 
     role: Required[Literal["system"]]

src/memos/types/openai_chat_completion_types/chat_completion_tool_message_param.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
 from typing import Literal
 
 from typing_extensions import Required, TypedDict
@@ -14,7 +13,7 @@
 
 
 class ChatCompletionToolMessageParam(TypedDict, total=False):
-    content: Required[str | Iterable[ChatCompletionContentPartParam]]
+    content: Required[str | list[ChatCompletionContentPartParam] | ChatCompletionContentPartParam]
     """The contents of the tool message."""
 
     role: Required[Literal["tool"]]

src/memos/types/openai_chat_completion_types/chat_completion_user_message_param.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
 from typing import Literal
 
 from typing_extensions import Required, TypedDict
@@ -14,7 +13,7 @@
 
 
 class ChatCompletionUserMessageParam(TypedDict, total=False):
-    content: Required[str | Iterable[ChatCompletionContentPartParam]]
+    content: Required[str | list[ChatCompletionContentPartParam] | ChatCompletionContentPartParam]
     """The contents of the user message."""
 
     role: Required[Literal["user"]]