Source code for langchain_community.document_loaders.pubmed

from typing import Iterator, Optional

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader
from langchain_community.utilities.pubmed import PubMedAPIWrapper


class PubMedLoader(BaseLoader):
    """从`PubMed`生物医学库加载。

    属性：
        query: 要传递给PubMed API的查询。
        load_max_docs: 要加载的最大文档数。"""

    def __init__(
        self,
        query: str,
        load_max_docs: Optional[int] = 3,
    ):
        """初始化PubMedLoader。

参数：
    query：要传递给PubMed API的查询。
    load_max_docs：要加载的最大文档数。默认为3。
"""
        self.query = query
        self.load_max_docs = load_max_docs
        self._client = PubMedAPIWrapper(  # type: ignore[call-arg]
            top_k_results=load_max_docs,  # type: ignore[arg-type]
        )

    def lazy_load(self) -> Iterator[Document]:
        for doc in self._client.lazy_load_docs(self.query):
            yield doc