diff --git a/libs/community/langchain_community/document_loaders/cube_semantic.py b/libs/community/langchain_community/document_loaders/cube_semantic.py index 26ebc5e2637..81a35787edd 100644 --- a/libs/community/langchain_community/document_loaders/cube_semantic.py +++ b/libs/community/langchain_community/document_loaders/cube_semantic.py @@ -1,7 +1,7 @@ import json import logging import time -from typing import List +from typing import Iterator, List import requests from langchain_core.documents import Document @@ -99,7 +99,7 @@ class CubeSemanticLoader(BaseLoader): logger.info("Maximum retries reached.") return [] - def load(self) -> List[Document]: + def lazy_load(self) -> Iterator[Document]: """Makes a call to Cube's REST API metadata endpoint. Returns: @@ -131,8 +131,6 @@ class CubeSemanticLoader(BaseLoader): if not cube_data_objects: raise ValueError("No cubes found in metadata.") - docs = [] - for cube_data_obj in cube_data_objects: cube_data_obj_name = cube_data_obj.get("name") cube_data_obj_type = cube_data_obj.get("type") @@ -173,6 +171,4 @@ class CubeSemanticLoader(BaseLoader): page_content = f"{str(item.get('title'))}, " page_content += f"{str(item.get('description'))}" - docs.append(Document(page_content=page_content, metadata=metadata)) - - return docs + yield Document(page_content=page_content, metadata=metadata)