diff --git a/src/chat_processor.py b/src/chat_processor.py index 75e4c698c..4fddae0b7 100644 --- a/src/chat_processor.py +++ b/src/chat_processor.py @@ -280,10 +280,54 @@ class ChatProcessor: web_sources = [] if use_web: try: - web_context, web_sources = comprehensive_web_search( - message, time_filter=time_filter, return_sources=True - ) - preface.append(untrusted_context_message("web search results", web_context)) + from src.llm_core import llm_call + + t_url, t_model, t_headers = session.endpoint_url, session.model, session.headers + + # Default fallback is the first non-empty line of the original user message + fallback_query = next((line.strip() for line in message.split("\n") if line.strip()), "") + search_query = fallback_query + + try: + generated_query = llm_call( + t_url, + t_model, + [ + { + "role": "system", + "content": ( + "Extract a concise search query from the user's message. " + "Reply ONLY with the query." + ), + }, + {"role": "user", "content": message}, + ], + headers=t_headers, + temperature=0.1, + max_tokens=50, + timeout=15, + ).strip() + + if generated_query: + # LLM successfully generated a non-empty query -> use the generated query + search_query = generated_query + else: + # LLM returned an empty or whitespace-only query -> fall back to original query + logger.warning("LLM generated an empty search query, using fallback.") + except Exception as e: + # LLM failed (exception/error) -> fall back to original user query + logger.warning(f"Failed to generate search query via LLM, using fallback: {e}") + + search_query = " ".join(search_query.split()) + if len(search_query) > 150: + search_query = search_query[:150].strip() + + if search_query: + # Execute web search using the final selected query + web_context, web_sources = comprehensive_web_search( + search_query, time_filter=time_filter, return_sources=True + ) + preface.append(untrusted_context_message("web search results", web_context)) except Exception as e: logger.error(f"Web search failed: {e}") preface.append({"role": "system", "content": "Web search encountered an error and could not retrieve results."}) diff --git a/tests/test_chat_processor_web_search.py b/tests/test_chat_processor_web_search.py new file mode 100644 index 000000000..06a271084 --- /dev/null +++ b/tests/test_chat_processor_web_search.py @@ -0,0 +1,95 @@ +from unittest.mock import MagicMock +from types import SimpleNamespace +from src.chat_processor import ChatProcessor + +def test_build_context_preface_web_search_success(monkeypatch): + """Test that LLM correctly extracts and uses a web search query.""" + mock_llm_call = MagicMock(return_value="extracted query") + monkeypatch.setattr("src.llm_core.llm_call", mock_llm_call) + + mock_web_search = MagicMock(return_value=("Search Results", [{"url": "http://mock.com"}])) + monkeypatch.setattr("src.chat_processor.comprehensive_web_search", mock_web_search) + + processor = ChatProcessor(memory_manager=MagicMock(), personal_docs_manager=MagicMock()) + session = SimpleNamespace(endpoint_url="http://local", model="test", headers={}) + + processor.build_context_preface( + message="Some text.\n\nSearch for LLMs.", + session=session, + use_web=True, + use_rag=False, + use_memory=False, + use_skills=False + ) + + mock_web_search.assert_called_with("extracted query", time_filter=None, return_sources=True) + +def test_build_context_preface_web_search_fallback_on_llm_failure(monkeypatch): + """Test fallback to original query if LLM fails.""" + def failing_llm(*args, **kwargs): + raise ValueError("LLM down") + monkeypatch.setattr("src.llm_core.llm_call", failing_llm) + + mock_web_search = MagicMock(return_value=("Search Results", [])) + monkeypatch.setattr("src.chat_processor.comprehensive_web_search", mock_web_search) + + processor = ChatProcessor(memory_manager=MagicMock(), personal_docs_manager=MagicMock()) + session = SimpleNamespace(endpoint_url="http://local", model="test", headers={}) + + processor.build_context_preface( + message="First line\nSecond line", + session=session, + use_web=True, + use_rag=False, + use_memory=False, + use_skills=False + ) + + mock_web_search.assert_called_with("First line", time_filter=None, return_sources=True) + +def test_build_context_preface_web_search_fallback_on_empty_generation(monkeypatch): + """Test fallback to original query if LLM returns empty string.""" + mock_llm_call = MagicMock(return_value=" \n ") + monkeypatch.setattr("src.llm_core.llm_call", mock_llm_call) + + mock_web_search = MagicMock(return_value=("Search Results", [])) + monkeypatch.setattr("src.chat_processor.comprehensive_web_search", mock_web_search) + + processor = ChatProcessor(memory_manager=MagicMock(), personal_docs_manager=MagicMock()) + session = SimpleNamespace(endpoint_url="http://local", model="test", headers={}) + + processor.build_context_preface( + message="\n\nFallback line\nNext", + session=session, + use_web=True, + use_rag=False, + use_memory=False, + use_skills=False + ) + + mock_web_search.assert_called_with("Fallback line", time_filter=None, return_sources=True) + +def test_build_context_preface_web_search_query_sanitization(monkeypatch): + """Test that query is truncated and whitespace collapsed.""" + long_query = "word " * 50 + mock_llm_call = MagicMock(return_value=long_query) + monkeypatch.setattr("src.llm_core.llm_call", mock_llm_call) + + mock_web_search = MagicMock(return_value=("Search Results", [])) + monkeypatch.setattr("src.chat_processor.comprehensive_web_search", mock_web_search) + + processor = ChatProcessor(memory_manager=MagicMock(), personal_docs_manager=MagicMock()) + session = SimpleNamespace(endpoint_url="http://local", model="test", headers={}) + + processor.build_context_preface( + message="Message", + session=session, + use_web=True, + use_rag=False, + use_memory=False, + use_skills=False + ) + + called_query = mock_web_search.call_args[0][0] + assert len(called_query) <= 150 + assert " " not in called_query