codekingpro/portable-devtools

1from typing import Iterator, Optional2 3from langchain_core.documents import Document4 5from langchain_community.document_loaders.base import BaseLoader6from langchain_community.utilities.pubmed import PubMedAPIWrapper7 8 9class PubMedLoader(BaseLoader):10    """Load from the `PubMed` biomedical library.11 12    Attributes:13        query: The query to be passed to the PubMed API.14        load_max_docs: The maximum number of documents to load.15    """16 17    def __init__(18        self,19        query: str,20        load_max_docs: Optional[int] = 3,21    ):22        """Initialize the PubMedLoader.23 24        Args:25            query: The query to be passed to the PubMed API.26            load_max_docs: The maximum number of documents to load.27              Defaults to 3.28        """29        self.query = query30        self.load_max_docs = load_max_docs31        self._client = PubMedAPIWrapper(  # type: ignore[call-arg]32            top_k_results=load_max_docs,  # type: ignore[arg-type]33        )34 35    def lazy_load(self) -> Iterator[Document]:36        for doc in self._client.lazy_load_docs(self.query):37            yield doc38