SearchArena / brave_ai.py
LeoWalker's picture
Results has a better user interface now
c7aea9b
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