Extract internet search results, chunk, and select most relevant chunks for LLM answer (INCOMPLETE)

Author: trial-danswer

backend/onyx/chat/process_message.py

@@ -43,6 +43,9 @@
 from onyx.chat.prompt_builder.answer_prompt_builder import AnswerPromptBuilder
 from onyx.chat.prompt_builder.answer_prompt_builder import default_build_system_message
 from onyx.chat.prompt_builder.answer_prompt_builder import default_build_user_message
+from onyx.chat.prompt_builder.citations_prompt import (
+    compute_max_document_tokens_for_persona,
+)
 from onyx.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
 from onyx.configs.chat_configs import DISABLE_LLM_CHOOSE_SEARCH
 from onyx.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
@@ -92,6 +95,7 @@
 from onyx.db.models import UserFile
 from onyx.db.persona import get_persona_by_id
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.user_documents import calculate_user_files_token_count
 from onyx.document_index.factory import get_default_document_index
 from onyx.file_store.models import ChatFileType
 from onyx.file_store.models import FileDescriptor
@@ -133,13 +137,13 @@
     INTERNET_SEARCH_RESPONSE_ID,
 )
 from onyx.tools.tool_implementations.internet_search.internet_search_tool import (
-    internet_search_response_to_search_docs,
+    InternetSearchTool,
 )
-from onyx.tools.tool_implementations.internet_search.internet_search_tool import (
+from onyx.tools.tool_implementations.internet_search.models import (
     InternetSearchResponse,
 )
-from onyx.tools.tool_implementations.internet_search.internet_search_tool import (
-    InternetSearchTool,
+from onyx.tools.tool_implementations.internet_search.utils import (
+    internet_search_response_to_search_docs,
 )
 from onyx.tools.tool_implementations.search.search_tool import (
     FINAL_CONTEXT_DOCUMENTS_ID,
@@ -279,6 +283,9 @@ def _handle_search_tool_response_summary(
     )
 
 
+# TODO: this takes the entire internet search response and sends it to LLM --> not correct
+# TODO: Internet search yields first an InternetSearchResponse to populate search results
+# and then yields a list of LlmDocs that should be added to context
 def _handle_internet_search_tool_response_summary(
     packet: ToolResponse,
     db_session: Session,
@@ -864,10 +871,6 @@ def stream_chat_message_objects(
                 file_id_to_user_file = {file.file_id: file for file in user_files}
 
             # Calculate token count for the files
-            from onyx.db.user_documents import calculate_user_files_token_count
-            from onyx.chat.prompt_builder.citations_prompt import (
-                compute_max_document_tokens_for_persona,
-            )
 
             total_tokens = calculate_user_files_token_count(
                 user_file_ids or [],
@@ -1048,6 +1051,16 @@ def create_response(
             structured_response_format=new_msg_req.structured_response_format,
         )
 
+        # Temp to get a pruning config for internet search
+        available_tokens = compute_max_document_tokens_for_persona(
+            db_session=db_session,
+            persona=persona,
+            actual_user_input=message_text,
+        )
+
+        internet_pruning_config = document_pruning_config.copy()
+        internet_pruning_config.max_tokens = available_tokens
+
         tool_dict = construct_tools(
             persona=persona,
             prompt_config=prompt_config,
@@ -1071,6 +1084,7 @@ def create_response(
             ),
             internet_search_tool_config=InternetSearchToolConfig(
                 answer_style_config=answer_style_config,
+                document_pruning_config=internet_pruning_config,
             ),
             image_generation_tool_config=ImageGenerationToolConfig(
                 additional_headers=litellm_additional_headers,

backend/onyx/tools/tool_constructor.py

@@ -122,6 +122,9 @@ class InternetSearchToolConfig(BaseModel):
             citation_config=CitationConfig(all_docs_useful=True)
         )
     )
+    document_pruning_config: DocumentPruningConfig = Field(
+        default_factory=DocumentPruningConfig
+    )
 
 
 class ImageGenerationToolConfig(BaseModel):
@@ -222,8 +225,11 @@ def construct_tools(
                 tool_dict[db_tool_model.id] = [
                     InternetSearchTool(
                         api_key=EXA_API_KEY,
+                        db_session=db_session,
+                        llm=llm,
                         answer_style_config=internet_search_tool_config.answer_style_config,
                         prompt_config=prompt_config,
+                        pruning_config=internet_search_tool_config.document_pruning_config,
                     )
                 ]