mirror of
https://github.com/open-webui/open-webui.git
synced 2026-09-16 23:43:03 +00:00
feat(retrieval): add Brave Search configuration options and update search function
This commit is contained in:
Yash
parent
0da58d7d16
commit
636e13eb93
4 changed files with 161 additions and 7 deletions
|
|
@ -3310,6 +3310,61 @@ BRAVE_SEARCH_API_KEY = PersistentConfig(
|
|||
os.getenv('BRAVE_SEARCH_API_KEY', ''),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_COUNTRY = PersistentConfig(
|
||||
'BRAVE_SEARCH_COUNTRY',
|
||||
'rag.web.search.brave_search_country',
|
||||
os.getenv('BRAVE_SEARCH_COUNTRY', ''),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_LANG = PersistentConfig(
|
||||
'BRAVE_SEARCH_LANG',
|
||||
'rag.web.search.brave_search_lang',
|
||||
os.getenv('BRAVE_SEARCH_LANG', ''),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_SPELLCHECK = PersistentConfig(
|
||||
'BRAVE_SEARCH_SPELLCHECK',
|
||||
'rag.web.search.brave_search_spellcheck',
|
||||
(os.getenv('BRAVE_SEARCH_SPELLCHECK', '').lower() == 'true' if os.getenv('BRAVE_SEARCH_SPELLCHECK') else None),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_MAX_TOKENS = PersistentConfig(
|
||||
'BRAVE_SEARCH_MAX_TOKENS',
|
||||
'rag.web.search.brave_search_max_tokens',
|
||||
int(os.getenv('BRAVE_SEARCH_MAX_TOKENS', '8192')),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_MAX_TOKENS_PER_URL = PersistentConfig(
|
||||
'BRAVE_SEARCH_MAX_TOKENS_PER_URL',
|
||||
'rag.web.search.brave_search_max_tokens_per_url',
|
||||
int(os.getenv('BRAVE_SEARCH_MAX_TOKENS_PER_URL', '4096')),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_MAX_SNIPPETS_PER_URL = PersistentConfig(
|
||||
'BRAVE_SEARCH_MAX_SNIPPETS_PER_URL',
|
||||
'rag.web.search.brave_search_max_snippets_per_url',
|
||||
int(os.getenv('BRAVE_SEARCH_MAX_SNIPPETS_PER_URL', '50')),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE = PersistentConfig(
|
||||
'BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE',
|
||||
'rag.web.search.brave_search_context_threshold_mode',
|
||||
os.getenv('BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE', 'balanced'),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_FRESHNESS = PersistentConfig(
|
||||
'BRAVE_SEARCH_FRESHNESS',
|
||||
'rag.web.search.brave_search_freshness',
|
||||
os.getenv('BRAVE_SEARCH_FRESHNESS', ''),
|
||||
)
|
||||
|
||||
BRAVE_SEARCH_GOGGLES = PersistentConfig(
|
||||
'BRAVE_SEARCH_GOGGLES',
|
||||
'rag.web.search.brave_search_goggles',
|
||||
os.getenv('BRAVE_SEARCH_GOGGLES', ''),
|
||||
)
|
||||
|
||||
|
||||
KAGI_SEARCH_API_KEY = PersistentConfig(
|
||||
'KAGI_SEARCH_API_KEY',
|
||||
'rag.web.search.kagi_search_api_key',
|
||||
|
|
|
|||
|
|
@ -344,6 +344,15 @@ from open_webui.config import (
|
|||
BING_SEARCH_V7_ENDPOINT,
|
||||
BING_SEARCH_V7_SUBSCRIPTION_KEY,
|
||||
BRAVE_SEARCH_API_KEY,
|
||||
BRAVE_SEARCH_COUNTRY,
|
||||
BRAVE_SEARCH_LANG,
|
||||
BRAVE_SEARCH_SPELLCHECK,
|
||||
BRAVE_SEARCH_MAX_TOKENS,
|
||||
BRAVE_SEARCH_MAX_TOKENS_PER_URL,
|
||||
BRAVE_SEARCH_MAX_SNIPPETS_PER_URL,
|
||||
BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE,
|
||||
BRAVE_SEARCH_FRESHNESS,
|
||||
BRAVE_SEARCH_GOGGLES,
|
||||
EXA_API_KEY,
|
||||
PERPLEXITY_API_KEY,
|
||||
PERPLEXITY_MODEL,
|
||||
|
|
@ -1102,6 +1111,15 @@ app.state.config.YACY_PASSWORD = YACY_PASSWORD
|
|||
app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
|
||||
app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
|
||||
app.state.config.BRAVE_SEARCH_API_KEY = BRAVE_SEARCH_API_KEY
|
||||
app.state.config.BRAVE_SEARCH_COUNTRY = BRAVE_SEARCH_COUNTRY
|
||||
app.state.config.BRAVE_SEARCH_LANG = BRAVE_SEARCH_LANG
|
||||
app.state.config.BRAVE_SEARCH_SPELLCHECK = BRAVE_SEARCH_SPELLCHECK
|
||||
app.state.config.BRAVE_SEARCH_MAX_TOKENS = BRAVE_SEARCH_MAX_TOKENS
|
||||
app.state.config.BRAVE_SEARCH_MAX_TOKENS_PER_URL = BRAVE_SEARCH_MAX_TOKENS_PER_URL
|
||||
app.state.config.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL = BRAVE_SEARCH_MAX_SNIPPETS_PER_URL
|
||||
app.state.config.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE = BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE
|
||||
app.state.config.BRAVE_SEARCH_FRESHNESS = BRAVE_SEARCH_FRESHNESS
|
||||
app.state.config.BRAVE_SEARCH_GOGGLES = BRAVE_SEARCH_GOGGLES
|
||||
app.state.config.KAGI_SEARCH_API_KEY = KAGI_SEARCH_API_KEY
|
||||
app.state.config.MOJEEK_SEARCH_API_KEY = MOJEEK_SEARCH_API_KEY
|
||||
app.state.config.BOCHA_SEARCH_API_KEY = BOCHA_SEARCH_API_KEY
|
||||
|
|
|
|||
|
|
@ -8,12 +8,37 @@ from open_webui.retrieval.web.main import SearchResult, get_filtered_results
|
|||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def search_brave(api_key: str, query: str, count: int, filter_list: Optional[list[str]] = None) -> list[SearchResult]:
|
||||
"""Search using Brave's Search API and return the results as a list of SearchResult objects.
|
||||
def search_brave(
|
||||
api_key: str,
|
||||
query: str,
|
||||
count: int,
|
||||
filter_list: Optional[list[str]] = None,
|
||||
country: Optional[str] = None,
|
||||
search_lang: Optional[str] = None,
|
||||
spellcheck: Optional[bool] = None,
|
||||
maximum_number_of_tokens: int = 8192,
|
||||
maximum_number_of_tokens_per_url: int = 4096,
|
||||
maximum_number_of_snippets_per_url: int = 50,
|
||||
context_threshold_mode: str = 'balanced',
|
||||
freshness: Optional[str] = None,
|
||||
goggles: Optional[str] = None,
|
||||
) -> list[SearchResult]:
|
||||
"""Search using Brave's LLM Context API and return the results as a list of SearchResult objects.
|
||||
|
||||
Args:
|
||||
api_key (str): A Brave Search API key
|
||||
query (str): The query to search for
|
||||
count (int): Maximum number of URLs to include in the context
|
||||
filter_list (list[str], optional): Domain allow/block list
|
||||
country (str, optional): 2-letter country code to localise results
|
||||
search_lang (str, optional): Language code for search results
|
||||
spellcheck (bool, optional): Whether to enable query spellcheck
|
||||
maximum_number_of_tokens (int): Approximate max tokens in context (default 8192, max 32768)
|
||||
maximum_number_of_tokens_per_url (int): Max tokens per URL (default 4096, max 8192)
|
||||
maximum_number_of_snippets_per_url (int): Max snippets per URL (default 50, max 100)
|
||||
context_threshold_mode (str): Threshold mode — 'balanced' or 'aggressive' (default 'balanced')
|
||||
freshness (str, optional): Age filter — pd/pw/pm/py or YYYY-MM-DDtoYYYY-MM-DD
|
||||
goggles (str, optional): Goggle URL or definition to rerank results
|
||||
"""
|
||||
|
||||
url = 'https://api.search.brave.com/res/v1/llm/context'
|
||||
|
|
@ -22,16 +47,27 @@ def search_brave(api_key: str, query: str, count: int, filter_list: Optional[lis
|
|||
'Accept-Encoding': 'gzip',
|
||||
'X-Subscription-Token': api_key,
|
||||
}
|
||||
|
||||
params = {
|
||||
'q': query,
|
||||
'maximum_number_of_urls': count,
|
||||
'maximum_number_of_tokens': 8192,
|
||||
'maximum_number_of_tokens_per_url': 4096,
|
||||
'maximum_number_of_snippets_per_url': 50,
|
||||
'context_threshold_mode': 'balanced',
|
||||
'maximum_number_of_tokens': maximum_number_of_tokens,
|
||||
'maximum_number_of_tokens_per_url': maximum_number_of_tokens_per_url,
|
||||
'maximum_number_of_snippets_per_url': maximum_number_of_snippets_per_url,
|
||||
'context_threshold_mode': context_threshold_mode,
|
||||
}
|
||||
|
||||
# Only include optional params when explicitly set
|
||||
if country:
|
||||
params['country'] = country
|
||||
if search_lang:
|
||||
params['search_lang'] = search_lang
|
||||
if spellcheck is not None:
|
||||
params['spellcheck'] = spellcheck
|
||||
if freshness:
|
||||
params['freshness'] = freshness
|
||||
if goggles:
|
||||
params['goggles'] = goggles
|
||||
|
||||
response = requests.get(url, headers=headers, params=params)
|
||||
|
||||
# Handle 429 rate limiting - Brave free tier allows 1 request/second
|
||||
|
|
|
|||
|
|
@ -531,6 +531,15 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
|
|||
'GOOGLE_PSE_API_KEY': request.app.state.config.GOOGLE_PSE_API_KEY,
|
||||
'GOOGLE_PSE_ENGINE_ID': request.app.state.config.GOOGLE_PSE_ENGINE_ID,
|
||||
'BRAVE_SEARCH_API_KEY': request.app.state.config.BRAVE_SEARCH_API_KEY,
|
||||
'BRAVE_SEARCH_COUNTRY': request.app.state.config.BRAVE_SEARCH_COUNTRY,
|
||||
'BRAVE_SEARCH_LANG': request.app.state.config.BRAVE_SEARCH_LANG,
|
||||
'BRAVE_SEARCH_SPELLCHECK': request.app.state.config.BRAVE_SEARCH_SPELLCHECK,
|
||||
'BRAVE_SEARCH_MAX_TOKENS': request.app.state.config.BRAVE_SEARCH_MAX_TOKENS,
|
||||
'BRAVE_SEARCH_MAX_TOKENS_PER_URL': request.app.state.config.BRAVE_SEARCH_MAX_TOKENS_PER_URL,
|
||||
'BRAVE_SEARCH_MAX_SNIPPETS_PER_URL': request.app.state.config.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL,
|
||||
'BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE': request.app.state.config.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE,
|
||||
'BRAVE_SEARCH_FRESHNESS': request.app.state.config.BRAVE_SEARCH_FRESHNESS,
|
||||
'BRAVE_SEARCH_GOGGLES': request.app.state.config.BRAVE_SEARCH_GOGGLES,
|
||||
'KAGI_SEARCH_API_KEY': request.app.state.config.KAGI_SEARCH_API_KEY,
|
||||
'MOJEEK_SEARCH_API_KEY': request.app.state.config.MOJEEK_SEARCH_API_KEY,
|
||||
'BOCHA_SEARCH_API_KEY': request.app.state.config.BOCHA_SEARCH_API_KEY,
|
||||
|
|
@ -599,6 +608,15 @@ class WebConfig(BaseModel):
|
|||
GOOGLE_PSE_API_KEY: Optional[str] = None
|
||||
GOOGLE_PSE_ENGINE_ID: Optional[str] = None
|
||||
BRAVE_SEARCH_API_KEY: Optional[str] = None
|
||||
BRAVE_SEARCH_COUNTRY: Optional[str] = None
|
||||
BRAVE_SEARCH_LANG: Optional[str] = None
|
||||
BRAVE_SEARCH_SPELLCHECK: Optional[bool] = None
|
||||
BRAVE_SEARCH_MAX_TOKENS: Optional[int] = None
|
||||
BRAVE_SEARCH_MAX_TOKENS_PER_URL: Optional[int] = None
|
||||
BRAVE_SEARCH_MAX_SNIPPETS_PER_URL: Optional[int] = None
|
||||
BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE: Optional[str] = None
|
||||
BRAVE_SEARCH_FRESHNESS: Optional[str] = None
|
||||
BRAVE_SEARCH_GOGGLES: Optional[str] = None
|
||||
KAGI_SEARCH_API_KEY: Optional[str] = None
|
||||
MOJEEK_SEARCH_API_KEY: Optional[str] = None
|
||||
BOCHA_SEARCH_API_KEY: Optional[str] = None
|
||||
|
|
@ -1081,6 +1099,15 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
|||
request.app.state.config.GOOGLE_PSE_API_KEY = form_data.web.GOOGLE_PSE_API_KEY
|
||||
request.app.state.config.GOOGLE_PSE_ENGINE_ID = form_data.web.GOOGLE_PSE_ENGINE_ID
|
||||
request.app.state.config.BRAVE_SEARCH_API_KEY = form_data.web.BRAVE_SEARCH_API_KEY
|
||||
request.app.state.config.BRAVE_SEARCH_COUNTRY = form_data.web.BRAVE_SEARCH_COUNTRY
|
||||
request.app.state.config.BRAVE_SEARCH_LANG = form_data.web.BRAVE_SEARCH_LANG
|
||||
request.app.state.config.BRAVE_SEARCH_SPELLCHECK = form_data.web.BRAVE_SEARCH_SPELLCHECK
|
||||
request.app.state.config.BRAVE_SEARCH_MAX_TOKENS = form_data.web.BRAVE_SEARCH_MAX_TOKENS
|
||||
request.app.state.config.BRAVE_SEARCH_MAX_TOKENS_PER_URL = form_data.web.BRAVE_SEARCH_MAX_TOKENS_PER_URL
|
||||
request.app.state.config.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL = form_data.web.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL
|
||||
request.app.state.config.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE = form_data.web.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE
|
||||
request.app.state.config.BRAVE_SEARCH_FRESHNESS = form_data.web.BRAVE_SEARCH_FRESHNESS
|
||||
request.app.state.config.BRAVE_SEARCH_GOGGLES = form_data.web.BRAVE_SEARCH_GOGGLES
|
||||
request.app.state.config.KAGI_SEARCH_API_KEY = form_data.web.KAGI_SEARCH_API_KEY
|
||||
request.app.state.config.MOJEEK_SEARCH_API_KEY = form_data.web.MOJEEK_SEARCH_API_KEY
|
||||
request.app.state.config.BOCHA_SEARCH_API_KEY = form_data.web.BOCHA_SEARCH_API_KEY
|
||||
|
|
@ -1216,6 +1243,15 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
|||
'GOOGLE_PSE_API_KEY': request.app.state.config.GOOGLE_PSE_API_KEY,
|
||||
'GOOGLE_PSE_ENGINE_ID': request.app.state.config.GOOGLE_PSE_ENGINE_ID,
|
||||
'BRAVE_SEARCH_API_KEY': request.app.state.config.BRAVE_SEARCH_API_KEY,
|
||||
'BRAVE_SEARCH_COUNTRY': request.app.state.config.BRAVE_SEARCH_COUNTRY,
|
||||
'BRAVE_SEARCH_LANG': request.app.state.config.BRAVE_SEARCH_LANG,
|
||||
'BRAVE_SEARCH_SPELLCHECK': request.app.state.config.BRAVE_SEARCH_SPELLCHECK,
|
||||
'BRAVE_SEARCH_MAX_TOKENS': request.app.state.config.BRAVE_SEARCH_MAX_TOKENS,
|
||||
'BRAVE_SEARCH_MAX_TOKENS_PER_URL': request.app.state.config.BRAVE_SEARCH_MAX_TOKENS_PER_URL,
|
||||
'BRAVE_SEARCH_MAX_SNIPPETS_PER_URL': request.app.state.config.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL,
|
||||
'BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE': request.app.state.config.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE,
|
||||
'BRAVE_SEARCH_FRESHNESS': request.app.state.config.BRAVE_SEARCH_FRESHNESS,
|
||||
'BRAVE_SEARCH_GOGGLES': request.app.state.config.BRAVE_SEARCH_GOGGLES,
|
||||
'KAGI_SEARCH_API_KEY': request.app.state.config.KAGI_SEARCH_API_KEY,
|
||||
'MOJEEK_SEARCH_API_KEY': request.app.state.config.MOJEEK_SEARCH_API_KEY,
|
||||
'BOCHA_SEARCH_API_KEY': request.app.state.config.BOCHA_SEARCH_API_KEY,
|
||||
|
|
@ -1960,6 +1996,15 @@ def search_web(request: Request, engine: str, query: str, user=None) -> list[Sea
|
|||
query,
|
||||
request.app.state.config.WEB_SEARCH_RESULT_COUNT,
|
||||
request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
|
||||
country=request.app.state.config.BRAVE_SEARCH_COUNTRY or None,
|
||||
search_lang=request.app.state.config.BRAVE_SEARCH_LANG or None,
|
||||
spellcheck=request.app.state.config.BRAVE_SEARCH_SPELLCHECK,
|
||||
maximum_number_of_tokens=request.app.state.config.BRAVE_SEARCH_MAX_TOKENS,
|
||||
maximum_number_of_tokens_per_url=request.app.state.config.BRAVE_SEARCH_MAX_TOKENS_PER_URL,
|
||||
maximum_number_of_snippets_per_url=request.app.state.config.BRAVE_SEARCH_MAX_SNIPPETS_PER_URL,
|
||||
context_threshold_mode=request.app.state.config.BRAVE_SEARCH_CONTEXT_THRESHOLD_MODE,
|
||||
freshness=request.app.state.config.BRAVE_SEARCH_FRESHNESS or None,
|
||||
goggles=request.app.state.config.BRAVE_SEARCH_GOGGLES or None,
|
||||
)
|
||||
else:
|
||||
raise Exception('No BRAVE_SEARCH_API_KEY found in environment variables')
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue