diff --git a/modules/chat.py b/modules/chat.py index 498c0d88..b2aacd5c 100644 --- a/modules/chat.py +++ b/modules/chat.py @@ -31,6 +31,7 @@ from modules.text_generation import ( get_max_prompt_length ) from modules.utils import delete_file, get_available_characters, save_file +from modules.web_search import add_web_search_attachments def strftime_now(format): @@ -566,6 +567,9 @@ def chatbot_wrapper(text, state, regenerate=False, _continue=False, loading_mess for file_path in files: add_message_attachment(output, row_idx, file_path, is_user=True) + # Add web search results as attachments if enabled + add_web_search_attachments(output, row_idx, text, state) + # Apply extensions text, visible_text = apply_extensions('chat_input', text, visible_text, state) text = apply_extensions('input', text, state, is_chat=True) diff --git a/modules/ui.py b/modules/ui.py index 00393b53..e24e6402 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -157,8 +157,6 @@ def list_model_elements(): def list_interface_input_elements(): elements = [ - 'navigate_message_index', - 'navigate_direction', 'temperature', 'dynatemp_low', 'dynatemp_high', @@ -218,6 +216,10 @@ def list_interface_input_elements(): 'edit_message_text', 'edit_message_role', 'branch_index', + 'enable_web_search', + 'web_search_pages', + 'navigate_message_index', + 'navigate_direction', ] # Chat elements diff --git a/modules/ui_chat.py b/modules/ui_chat.py index 952a40a5..719af85a 100644 --- a/modules/ui_chat.py +++ b/modules/ui_chat.py @@ -86,6 +86,12 @@ def create_ui(): with gr.Row(): shared.gradio['start_with'] = gr.Textbox(label='Start reply with', placeholder='Sure thing!', value=shared.settings['start_with'], elem_classes=['add_scrollbar']) + with gr.Row(): + shared.gradio['enable_web_search'] = gr.Checkbox(value=shared.settings.get('enable_web_search', False), label='Activate web search') + + with gr.Row(visible=shared.settings.get('enable_web_search', False)) as shared.gradio['web_search_row']: + shared.gradio['web_search_pages'] = gr.Number(value=shared.settings.get('web_search_pages', 3), precision=0, label='Number of pages to download', minimum=1, maximum=10) + with gr.Row(): shared.gradio['mode'] = gr.Radio(choices=['instruct', 'chat-instruct', 'chat'], value=shared.settings['mode'] if shared.settings['mode'] in ['chat', 'chat-instruct'] else None, label='Mode', info='Defines how the chat prompt is generated. In instruct and chat-instruct modes, the instruction template Parameters > Instruction template is used.', elem_id='chat-mode') @@ -369,3 +375,9 @@ def create_event_handlers(): shared.gradio['count_tokens'].click( ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then( chat.count_prompt_tokens, gradio('textbox', 'interface_state'), gradio('token_display'), show_progress=False) + + shared.gradio['enable_web_search'].change( + lambda x: gr.update(visible=x), + gradio('enable_web_search'), + gradio('web_search_row') + ) diff --git a/modules/web_search.py b/modules/web_search.py new file mode 100644 index 00000000..e7688ba4 --- /dev/null +++ b/modules/web_search.py @@ -0,0 +1,125 @@ +from datetime import datetime + +import requests +from bs4 import BeautifulSoup +from duckduckgo_search import DDGS + +from modules.logging_colors import logger +from modules.text_generation import generate_reply + + +def get_current_timestamp(): + """Returns the current time in 24-hour format""" + return datetime.now().strftime('%b %d, %Y %H:%M') + + +def generate_search_query(user_message, state): + """Generate a search query from user message using the LLM""" + search_prompt = f"{user_message}\n\n=====\n\nPlease turn the message above into a short web search query in the same language as the message. Respond with only the search query, nothing else." + + # Use a minimal state for search query generation + search_state = state.copy() + search_state['max_new_tokens'] = 64 + search_state['temperature'] = 0.1 + + query = "" + for reply in generate_reply(search_prompt, search_state, stopping_strings=[], is_chat=False): + query = reply.strip() + + return query + + +def download_web_page(url, timeout=10): + """Download and extract text from a web page""" + try: + headers = { + 'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36' + } + response = requests.get(url, headers=headers, timeout=timeout) + response.raise_for_status() + + soup = BeautifulSoup(response.content, 'html.parser') + + # Remove script and style elements + for script in soup(["script", "style"]): + script.decompose() + + # Get text and clean it up + text = soup.get_text() + lines = (line.strip() for line in text.splitlines()) + chunks = (phrase.strip() for line in lines for phrase in line.split(" ")) + text = ' '.join(chunk for chunk in chunks if chunk) + + return text + except Exception as e: + logger.error(f"Error downloading {url}: {e}") + return f"[Error downloading content from {url}: {str(e)}]" + + +def perform_web_search(query, num_pages=3): + """Perform web search and return results with content""" + try: + with DDGS() as ddgs: + results = list(ddgs.text(query, max_results=num_pages)) + + search_results = [] + for i, result in enumerate(results): + url = result.get('href', '') + title = result.get('title', f'Search Result {i+1}') + + # Download page content + content = download_web_page(url) + + search_results.append({ + 'title': title, + 'url': url, + 'content': content + }) + + return search_results + except Exception as e: + logger.error(f"Error performing web search: {e}") + return [] + + +def add_web_search_attachments(history, row_idx, user_message, state): + """Perform web search and add results as attachments""" + if not state.get('enable_web_search', False): + return + + try: + # Generate search query + search_query = generate_search_query(user_message, state) + if not search_query: + logger.warning("Failed to generate search query") + return + + logger.info(f"Generated search query: {search_query}") + + # Perform web search + num_pages = int(state.get('web_search_pages', 3)) + search_results = perform_web_search(search_query, num_pages) + + if not search_results: + logger.warning("No search results found") + return + + # Add search results as attachments + key = f"user_{row_idx}" + if key not in history['metadata']: + history['metadata'][key] = {"timestamp": get_current_timestamp()} + if "attachments" not in history['metadata'][key]: + history['metadata'][key]["attachments"] = [] + + for result in search_results: + attachment = { + "name": f"{result['title']}", + "type": "text/html", + "content": f"URL: {result['url']}\n\n{result['content']}" + } + history['metadata'][key]["attachments"].append(attachment) + + logger.info(f"Added {len(search_results)} web search results as attachments") + + except Exception as e: + logger.error(f"Error in web search: {e}") diff --git a/requirements/full/requirements.txt b/requirements/full/requirements.txt index 3d18f5fd..0eaf10da 100644 --- a/requirements/full/requirements.txt +++ b/requirements/full/requirements.txt @@ -1,7 +1,9 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 bitsandbytes==0.45.* colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_amd.txt b/requirements/full/requirements_amd.txt index 82b19964..65f184bf 100644 --- a/requirements/full/requirements_amd.txt +++ b/requirements/full/requirements_amd.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_amd_noavx2.txt b/requirements/full/requirements_amd_noavx2.txt index a8b03014..d20b2ec3 100644 --- a/requirements/full/requirements_amd_noavx2.txt +++ b/requirements/full/requirements_amd_noavx2.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_apple_intel.txt b/requirements/full/requirements_apple_intel.txt index 5a61ac7d..2613d787 100644 --- a/requirements/full/requirements_apple_intel.txt +++ b/requirements/full/requirements_apple_intel.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_apple_silicon.txt b/requirements/full/requirements_apple_silicon.txt index 6862c3b4..af583b00 100644 --- a/requirements/full/requirements_apple_silicon.txt +++ b/requirements/full/requirements_apple_silicon.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_cpu_only.txt b/requirements/full/requirements_cpu_only.txt index e6982779..9bf2a37d 100644 --- a/requirements/full/requirements_cpu_only.txt +++ b/requirements/full/requirements_cpu_only.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_cpu_only_noavx2.txt b/requirements/full/requirements_cpu_only_noavx2.txt index 97bff786..1731448e 100644 --- a/requirements/full/requirements_cpu_only_noavx2.txt +++ b/requirements/full/requirements_cpu_only_noavx2.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_noavx2.txt b/requirements/full/requirements_noavx2.txt index 17c7e246..fc481a1a 100644 --- a/requirements/full/requirements_noavx2.txt +++ b/requirements/full/requirements_noavx2.txt @@ -1,7 +1,9 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 bitsandbytes==0.45.* colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/full/requirements_nowheels.txt b/requirements/full/requirements_nowheels.txt index 89b32caf..2ed8affa 100644 --- a/requirements/full/requirements_nowheels.txt +++ b/requirements/full/requirements_nowheels.txt @@ -1,6 +1,8 @@ accelerate==1.5.* +beautifulsoup4==4.13.4 colorama datasets +duckduckgo_search==8.0.2 einops fastapi==0.112.4 gradio==4.37.* diff --git a/requirements/portable/requirements.txt b/requirements/portable/requirements.txt index ec9bafc6..fdae681d 100644 --- a/requirements/portable/requirements.txt +++ b/requirements/portable/requirements.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_apple_intel.txt b/requirements/portable/requirements_apple_intel.txt index 025a737e..a58f39f7 100644 --- a/requirements/portable/requirements_apple_intel.txt +++ b/requirements/portable/requirements_apple_intel.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_apple_silicon.txt b/requirements/portable/requirements_apple_silicon.txt index 32644e87..91ea3a6d 100644 --- a/requirements/portable/requirements_apple_silicon.txt +++ b/requirements/portable/requirements_apple_silicon.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_cpu_only.txt b/requirements/portable/requirements_cpu_only.txt index bd5c1d9b..37e5aa40 100644 --- a/requirements/portable/requirements_cpu_only.txt +++ b/requirements/portable/requirements_cpu_only.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_cpu_only_noavx2.txt b/requirements/portable/requirements_cpu_only_noavx2.txt index 51f2b7d9..dcb2884b 100644 --- a/requirements/portable/requirements_cpu_only_noavx2.txt +++ b/requirements/portable/requirements_cpu_only_noavx2.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_noavx2.txt b/requirements/portable/requirements_noavx2.txt index aad6bf5a..8f1295bb 100644 --- a/requirements/portable/requirements_noavx2.txt +++ b/requirements/portable/requirements_noavx2.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_nowheels.txt b/requirements/portable/requirements_nowheels.txt index 4c055426..21805fe2 100644 --- a/requirements/portable/requirements_nowheels.txt +++ b/requirements/portable/requirements_nowheels.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_vulkan.txt b/requirements/portable/requirements_vulkan.txt index 3d98d1b0..858b4488 100644 --- a/requirements/portable/requirements_vulkan.txt +++ b/requirements/portable/requirements_vulkan.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6 diff --git a/requirements/portable/requirements_vulkan_noavx2.txt b/requirements/portable/requirements_vulkan_noavx2.txt index f954b8d2..569bae99 100644 --- a/requirements/portable/requirements_vulkan_noavx2.txt +++ b/requirements/portable/requirements_vulkan_noavx2.txt @@ -1,3 +1,5 @@ +beautifulsoup4==4.13.4 +duckduckgo_search==8.0.2 fastapi==0.112.4 gradio==4.37.* jinja2==3.1.6