add some consistent doc ids (#868)

run-llama · Jan 16, 2024 · 80d8b41 · 80d8b41
1 parent 7f9a92c
commit 80d8b41
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 2 deletions.
diff --git a/llama_hub/discord/base.py b/llama_hub/discord/base.py
@@ -74,6 +74,7 @@ async def on_ready(self) -> None:
         map(
             lambda msg: Document(
                 text=msg.content,
+                id_=msg.id,
                 metadata={
                     "message_id": msg.id,
                     "username": msg.author.name,

diff --git a/llama_hub/google_drive/base.py b/llama_hub/google_drive/base.py
@@ -308,6 +308,8 @@ def get_metadata(filename):
                     SimpleDirectoryReader = download_loader("SimpleDirectoryReader")
                 loader = SimpleDirectoryReader(temp_dir, file_metadata=get_metadata)
                 documents = loader.load_data()
+                for doc in documents:
+                    doc.id_ = doc.metadata.get("file id", doc.id_)
 
             return documents
         except Exception as e:

diff --git a/llama_hub/notion/base.py b/llama_hub/notion/base.py
@@ -165,11 +165,19 @@ def load_data(
             page_ids = self.query_database(database_id)
             for page_id in page_ids:
                 page_text = self.read_page(page_id)
-                docs.append(Document(text=page_text, extra_info={"page_id": page_id}))
+                docs.append(
+                    Document(
+                        text=page_text, id_=page_id, extra_info={"page_id": page_id}
+                    )
+                )
         else:
             for page_id in page_ids:
                 page_text = self.read_page(page_id)
-                docs.append(Document(text=page_text, extra_info={"page_id": page_id}))
+                docs.append(
+                    Document(
+                        text=page_text, id_=page_id, extra_info={"page_id": page_id}
+                    )
+                )
 
         return docs
 

diff --git a/llama_hub/s3/base.py b/llama_hub/s3/base.py
@@ -151,4 +151,9 @@ def load_data(self, custom_temp_subdir: str = None) -> List[Document]:
             os.makedirs(temp_dir, exist_ok=True)
             documents = self.load_s3_files_as_docs(temp_dir)
             shutil.rmtree(temp_dir)
+
+        if not self.filename_as_id:
+            for doc in documents:
+                doc.id_ = self.s3_endpoint_url + "_" + doc.metadata["filename"]
+
         return documents
diff --git a/llama_hub/slack/base.py b/llama_hub/slack/base.py
@@ -188,6 +188,7 @@ def load_data(
                 results.append(
                     Document(
                         text=message_content["text"],
+                        id_=channel_id,
                         extra_info={
                             "channel": channel_id,
                             "timestamp": timestamp,