From 9e6d4eeb9277476603499b663019f680ccc87a99 Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Tue, 16 Apr 2024 17:09:15 +0800
Subject: [PATCH] fix the return with wrong datatype of segment (#3525)

---
 api/core/docstore/dataset_docstore.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/api/core/docstore/dataset_docstore.py b/api/core/docstore/dataset_docstore.py
index 9a051fd4c..7567493b9 100644
--- a/api/core/docstore/dataset_docstore.py
+++ b/api/core/docstore/dataset_docstore.py
@@ -84,7 +84,7 @@ class DatasetDocumentStore:
             if not isinstance(doc, Document):
                 raise ValueError("doc must be a Document")
 
-            segment_document = self.get_document(doc_id=doc.metadata['doc_id'], raise_error=False)
+            segment_document = self.get_document_segment(doc_id=doc.metadata['doc_id'])
 
             # NOTE: doc could already exist in the store, but we overwrite it
             if not allow_update and segment_document: