Spaces:
Sleeping
Sleeping
import json | |
import requests | |
from typing import List, Optional | |
# from langchain_community.document_transformers import MarkdownifyTransformer | |
from langchain_core.documents import Document | |
from langchain_core.pydantic_v1 import BaseModel, Field | |
class BraveAIWrapper(BaseModel): | |
api_key: str = Field(..., description="API key for Brave Search") | |
base_search_url: str = Field("https://api.search.brave.com/res/v1/web/search", const=True) | |
base_summarize_url: str = Field("https://api.search.brave.com/res/v1/summarizer/search", const=True) | |
headers: dict = Field(default_factory=dict, description="HTTP headers for API requests") | |
class Config: | |
arbitrary_types_allowed = True | |
def __init__(self, **data): | |
super().__init__(**data) | |
self.headers = { | |
"X-Subscription-Token": self.api_key, | |
"Accept": "application/json", | |
"Accept-Encoding": "gzip", | |
} | |
def get_brave_results(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[dict]: | |
""" | |
Get search results from Brave Search. | |
Args: | |
query (str): The search query. | |
count (int): Number of results to return. | |
safe_search (str): Safe search filter (off, moderate, strict). | |
Returns: | |
Optional[dict]: JSON response from Brave Search API or None if an error occurs. | |
""" | |
params = { | |
"q": query, | |
"count": count, | |
"summary": True, | |
"safe_search": safe_search, | |
"extra_snippets": True, | |
} | |
try: | |
response = requests.get(self.base_search_url, headers=self.headers, params=params) | |
response.raise_for_status() | |
return response.json() | |
except requests.exceptions.RequestException as e: | |
print(f"Error: {e}") | |
return None | |
def summarize_results(self, summarizer_key: str) -> Optional[dict]: | |
""" | |
Summarize search results using Brave Summarizer. | |
Args: | |
summarizer_key (str): The key for the summarizer. | |
Returns: | |
Optional[dict]: JSON response from Brave Summarizer API or None if an error occurs. | |
""" | |
params = {"key": summarizer_key} | |
try: | |
response = requests.get(self.base_summarize_url, headers=self.headers, params=params) | |
response.raise_for_status() | |
return response.json() | |
except requests.exceptions.RequestException as e: | |
print(f"Error: {e}") | |
return None | |
def get_and_summarize(self, query: str, count: int = 3, safe_search: str = 'off') -> Optional[str]: | |
""" | |
Get and summarize search results from Brave Search. | |
Args: | |
query (str): The search query. | |
count (int): Number of results to return. | |
safe_search (str): Safe search filter (off, moderate, strict). | |
Returns: | |
Optional[str]: Summarized result or None if an error occurs. | |
""" | |
results = self.get_brave_results(query, count, safe_search) | |
if results and 'summarizer' in results: | |
summarizer_key = results['summarizer']['key'] | |
summary = self.summarize_results(summarizer_key) | |
if summary and 'summary' in summary and len(summary['summary']) > 0: | |
return summary['summary'][0]['data'] | |
return None | |