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

Covered by `test_cube_semantic.py`
This commit is contained in:
Christophe Bornet 2024-03-06 02:32:31 +01:00 committed by GitHub
parent a6b5d45e31
commit 7d6de96186
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -1,7 +1,7 @@
import json import json
import logging import logging
import time import time
from typing import List from typing import Iterator, List
import requests import requests
from langchain_core.documents import Document from langchain_core.documents import Document
@ -99,7 +99,7 @@ class CubeSemanticLoader(BaseLoader):
logger.info("Maximum retries reached.") logger.info("Maximum retries reached.")
return [] return []
def load(self) -> List[Document]: def lazy_load(self) -> Iterator[Document]:
"""Makes a call to Cube's REST API metadata endpoint. """Makes a call to Cube's REST API metadata endpoint.
Returns: Returns:
@ -131,8 +131,6 @@ class CubeSemanticLoader(BaseLoader):
if not cube_data_objects: if not cube_data_objects:
raise ValueError("No cubes found in metadata.") raise ValueError("No cubes found in metadata.")
docs = []
for cube_data_obj in cube_data_objects: for cube_data_obj in cube_data_objects:
cube_data_obj_name = cube_data_obj.get("name") cube_data_obj_name = cube_data_obj.get("name")
cube_data_obj_type = cube_data_obj.get("type") 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('title'))}, "
page_content += f"{str(item.get('description'))}" page_content += f"{str(item.get('description'))}"
docs.append(Document(page_content=page_content, metadata=metadata)) yield Document(page_content=page_content, metadata=metadata)
return docs