mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-21 02:17:12 +00:00
ecff055096
- **Description:**
- Make Brave Search Tool consistent with other tools and allow reading
its api key from `BRAVE_SEARCH_API_KEY` instead of having to pass the
api key manually (no breaking changes)
- Improve Brave Search Tool by storing api key in `SecretStr` instead of
plain `str`.
- Add unit test for `BraveSearchWrapper`
- Reflect the changes in the documentation
- **Issue:** N/A
- **Dependencies:** N/A
- **Twitter handle:** ivan_brko
35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
from typing import Iterator, List, Optional
|
|
|
|
from langchain_core.documents import Document
|
|
from pydantic import SecretStr
|
|
|
|
from langchain_community.document_loaders.base import BaseLoader
|
|
from langchain_community.utilities.brave_search import BraveSearchWrapper
|
|
|
|
|
|
class BraveSearchLoader(BaseLoader):
|
|
"""Load with `Brave Search` engine."""
|
|
|
|
def __init__(self, query: str, api_key: str, search_kwargs: Optional[dict] = None):
|
|
"""Initializes the BraveLoader.
|
|
|
|
Args:
|
|
query: The query to search for.
|
|
api_key: The API key to use.
|
|
search_kwargs: The search kwargs to use.
|
|
"""
|
|
self.query = query
|
|
self.api_key = api_key
|
|
self.search_kwargs = search_kwargs or {}
|
|
|
|
def load(self) -> List[Document]:
|
|
brave_client = BraveSearchWrapper(
|
|
api_key=SecretStr(self.api_key),
|
|
search_kwargs=self.search_kwargs,
|
|
)
|
|
return brave_client.download_documents(self.query)
|
|
|
|
def lazy_load(self) -> Iterator[Document]:
|
|
for doc in self.load():
|
|
yield doc
|