codekingpro/portable-devtools

1from typing import Iterator2 3from langchain_core.documents import Document4 5from langchain_community.document_loaders.base import BaseLoader6 7 8class OpenCityDataLoader(BaseLoader):9    """Load from `Open City`."""10 11    def __init__(self, city_id: str, dataset_id: str, limit: int):12        """Initialize with dataset_id.13        Example: https://dev.socrata.com/foundry/data.sfgov.org/vw6y-z8j614        e.g., city_id = data.sfgov.org15        e.g., dataset_id = vw6y-z8j616 17        Args:18            city_id: The Open City city identifier.19            dataset_id: The Open City dataset identifier.20            limit: The maximum number of documents to load.21        """22        self.city_id = city_id23        self.dataset_id = dataset_id24        self.limit = limit25 26    def lazy_load(self) -> Iterator[Document]:27        """Lazy load records."""28 29        from sodapy import Socrata30 31        client = Socrata(self.city_id, None)32        results = client.get(self.dataset_id, limit=self.limit)33        for record in results:34            yield Document(35                page_content=str(record),36                metadata={37                    "source": self.city_id + "_" + self.dataset_id,38                },39            )40