From ed36f9f6046b5cd571967591974268c33ad46abb Mon Sep 17 00:00:00 2001 From: Christophe Bornet Date: Wed, 6 Mar 2024 19:03:46 +0100 Subject: [PATCH] community: Implement lazy_load() for WhatsAppChatLoader (#18677) Integration test: `tests/integration_tests/document_loaders/test_whatsapp_chat.py` --- .../langchain_community/document_loaders/whatsapp_chat.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/libs/community/langchain_community/document_loaders/whatsapp_chat.py b/libs/community/langchain_community/document_loaders/whatsapp_chat.py index d8198dee1cf..decda7f3ff9 100644 --- a/libs/community/langchain_community/document_loaders/whatsapp_chat.py +++ b/libs/community/langchain_community/document_loaders/whatsapp_chat.py @@ -1,6 +1,6 @@ import re from pathlib import Path -from typing import List +from typing import Iterator from langchain_core.documents import Document @@ -19,8 +19,7 @@ class WhatsAppChatLoader(BaseLoader): """Initialize with path.""" self.file_path = path - def load(self) -> List[Document]: - """Load documents.""" + def lazy_load(self) -> Iterator[Document]: p = Path(self.file_path) text_content = "" @@ -62,4 +61,4 @@ class WhatsAppChatLoader(BaseLoader): metadata = {"source": str(p)} - return [Document(page_content=text_content, metadata=metadata)] + yield Document(page_content=text_content, metadata=metadata)