Fix/file source (#640)

* Fallback source_doc_id to file_ids in KB logs

* Refactor(scheduler): Use file_ids directly for source_doc_id in KB logs

* Refactor(scheduler): Safely access file_ids for KB log source_doc_id

---------

Co-authored-by: [email protected] <>

Author: glin93

src/memos/mem_scheduler/general_scheduler.py

@@ -516,8 +516,12 @@ def send_add_log_messages_to_cloud_env(
         """
         kb_log_content: list[dict] = []
         info = msg.info or {}
+
         # Process added items
         for item in prepared_add_items:
+            metadata = getattr(item, "metadata", None)
+            file_ids = getattr(metadata, "file_ids", None) if metadata else None
+            source_doc_id = file_ids[0] if isinstance(file_ids, list) and file_ids else None
             kb_log_content.append(
                 {
                     "log_source": "KNOWLEDGE_BASE_LOG",
@@ -526,13 +530,16 @@ def send_add_log_messages_to_cloud_env(
                     "memory_id": item.id,
                     "content": item.memory,
                     "original_content": None,
-                    "source_doc_id": getattr(item.metadata, "source_doc_id", None),
+                    "source_doc_id": source_doc_id,
                 }
             )
 
         # Process updated items
         for item_data in prepared_update_items_with_original:
             item = item_data["new_item"]
+            metadata = getattr(item, "metadata", None)
+            file_ids = getattr(metadata, "file_ids", None) if metadata else None
+            source_doc_id = file_ids[0] if isinstance(file_ids, list) and file_ids else None
             kb_log_content.append(
                 {
                     "log_source": "KNOWLEDGE_BASE_LOG",
@@ -541,7 +548,7 @@ def send_add_log_messages_to_cloud_env(
                     "memory_id": item.id,
                     "content": item.memory,
                     "original_content": item_data.get("original_content"),
-                    "source_doc_id": getattr(item.metadata, "source_doc_id", None),
+                    "source_doc_id": source_doc_id,
                 }
             )
 
@@ -888,6 +895,11 @@ def _process_memories_with_reader(
                         # New: Knowledge Base Logging (Cloud Service)
                         kb_log_content = []
                         for item in flattened_memories:
+                            metadata = getattr(item, "metadata", None)
+                            file_ids = getattr(metadata, "file_ids", None) if metadata else None
+                            source_doc_id = (
+                                file_ids[0] if isinstance(file_ids, list) and file_ids else None
+                            )
                             kb_log_content.append(
                                 {
                                     "log_source": "KNOWLEDGE_BASE_LOG",
@@ -898,7 +910,7 @@ def _process_memories_with_reader(
                                     "memory_id": item.id,
                                     "content": item.memory,
                                     "original_content": None,
-                                    "source_doc_id": getattr(item.metadata, "source_doc_id", None),
+                                    "source_doc_id": source_doc_id,
                                 }
                             )
                         if kb_log_content: