codekingpro/portable-devtools

1from typing import Iterator, List, Optional2 3from langchain_core.documents import Document4from pydantic import SecretStr5 6from langchain_community.document_loaders.base import BaseLoader7from langchain_community.utilities.brave_search import BraveSearchWrapper8 9 10class BraveSearchLoader(BaseLoader):11    """Load with `Brave Search` engine."""12 13    def __init__(self, query: str, api_key: str, search_kwargs: Optional[dict] = None):14        """Initializes the BraveLoader.15 16        Args:17            query: The query to search for.18            api_key: The API key to use.19            search_kwargs: The search kwargs to use.20        """21        self.query = query22        self.api_key = api_key23        self.search_kwargs = search_kwargs or {}24 25    def load(self) -> List[Document]:26        brave_client = BraveSearchWrapper(27            api_key=SecretStr(self.api_key),28            search_kwargs=self.search_kwargs,29        )30        return brave_client.download_documents(self.query)31 32    def lazy_load(self) -> Iterator[Document]:33        for doc in self.load():34            yield doc35