community[patch]: Implement lazy_load() for CubeSemanticLoader (#18535)

Covered by `test_cube_semantic.py`
pull/18548/head^2
Christophe Bornet 4 months ago committed by GitHub
parent a6b5d45e31
commit 7d6de96186
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -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)

Loading…
Cancel
Save