diff --git a/backend/onyx/agents/agent_search/basic/utils.py b/backend/onyx/agents/agent_search/basic/utils.py index 51641adbf76..9fbba12cac1 100644 --- a/backend/onyx/agents/agent_search/basic/utils.py +++ b/backend/onyx/agents/agent_search/basic/utils.py @@ -7,6 +7,7 @@ from onyx.agents.agent_search.shared_graph_utils.utils import write_custom_event from onyx.chat.models import LlmDoc +from onyx.chat.models import OnyxContext from onyx.chat.stream_processing.answer_response_handler import AnswerResponseHandler from onyx.chat.stream_processing.answer_response_handler import CitationResponseHandler from onyx.chat.stream_processing.answer_response_handler import ( @@ -23,7 +24,7 @@ def process_llm_stream( should_stream_answer: bool, writer: StreamWriter, final_search_results: list[LlmDoc] | None = None, - displayed_search_results: list[LlmDoc] | None = None, + displayed_search_results: list[OnyxContext] | list[LlmDoc] | None = None, ) -> AIMessageChunk: tool_call_chunk = AIMessageChunk(content="") diff --git a/backend/onyx/agents/agent_search/orchestration/nodes/basic_use_tool_response.py b/backend/onyx/agents/agent_search/orchestration/nodes/basic_use_tool_response.py index 164551746fc..6874aae9795 100644 --- a/backend/onyx/agents/agent_search/orchestration/nodes/basic_use_tool_response.py +++ b/backend/onyx/agents/agent_search/orchestration/nodes/basic_use_tool_response.py @@ -9,6 +9,7 @@ from onyx.agents.agent_search.basic.utils import process_llm_stream from onyx.agents.agent_search.models import GraphConfig from onyx.chat.models import LlmDoc +from onyx.chat.models import OnyxContexts from onyx.tools.tool_implementations.search.search_tool import ( SEARCH_DOC_CONTENT_ID, ) @@ -50,13 +51,11 @@ def basic_use_tool_response( if yield_item.id == FINAL_CONTEXT_DOCUMENTS_ID: final_search_results = cast(list[LlmDoc], yield_item.response) elif yield_item.id == SEARCH_DOC_CONTENT_ID: - search_contexts = yield_item.response.contexts + search_contexts = cast(OnyxContexts, yield_item.response).contexts for doc in search_contexts: if doc.document_id not in initial_search_results: initial_search_results.append(doc) - initial_search_results = cast(list[LlmDoc], initial_search_results) - new_tool_call_chunk = AIMessageChunk(content="") if not agent_config.behavior.skip_gen_ai_answer_generation: stream = llm.stream( @@ -70,7 +69,9 @@ def basic_use_tool_response( True, writer, final_search_results=final_search_results, - displayed_search_results=initial_search_results, + # when the search tool is called with specific doc ids, initial search + # results are not output. But, we still want i.e. citations to be processed. + displayed_search_results=initial_search_results or final_search_results, ) return BasicOutput(tool_call_chunk=new_tool_call_chunk) diff --git a/backend/onyx/chat/answer.py b/backend/onyx/chat/answer.py index 303f1e6a39d..8addfe5ba83 100644 --- a/backend/onyx/chat/answer.py +++ b/backend/onyx/chat/answer.py @@ -183,6 +183,7 @@ def citations_by_subquestion(self) -> dict[SubQuestionKey, list[CitationInfo]]: citations_by_subquestion: dict[ SubQuestionKey, list[CitationInfo] ] = defaultdict(list) + basic_subq_key = SubQuestionKey(level=BASIC_KEY[0], question_num=BASIC_KEY[1]) for packet in self.processed_streamed_output: if isinstance(packet, CitationInfo): if packet.level_question_num is not None and packet.level is not None: @@ -192,7 +193,7 @@ def citations_by_subquestion(self) -> dict[SubQuestionKey, list[CitationInfo]]: ) ].append(packet) elif packet.level is None: - citations_by_subquestion[BASIC_SQ_KEY].append(packet) + citations_by_subquestion[basic_subq_key].append(packet) return citations_by_subquestion def is_cancelled(self) -> bool: diff --git a/backend/onyx/chat/stream_processing/utils.py b/backend/onyx/chat/stream_processing/utils.py index 97aa4ce607b..bcbcad5700e 100644 --- a/backend/onyx/chat/stream_processing/utils.py +++ b/backend/onyx/chat/stream_processing/utils.py @@ -3,6 +3,7 @@ from pydantic import BaseModel from onyx.chat.models import LlmDoc +from onyx.chat.models import OnyxContext from onyx.context.search.models import InferenceChunk @@ -11,7 +12,7 @@ class DocumentIdOrderMapping(BaseModel): def map_document_id_order( - chunks: Sequence[InferenceChunk | LlmDoc], one_indexed: bool = True + chunks: Sequence[InferenceChunk | LlmDoc | OnyxContext], one_indexed: bool = True ) -> DocumentIdOrderMapping: order_mapping = {} current = 1 if one_indexed else 0 diff --git a/backend/onyx/server/query_and_chat/chat_backend.py b/backend/onyx/server/query_and_chat/chat_backend.py index 0511683abf3..839c79ac24e 100644 --- a/backend/onyx/server/query_and_chat/chat_backend.py +++ b/backend/onyx/server/query_and_chat/chat_backend.py @@ -743,7 +743,7 @@ def upload_files_for_chat( # to re-extract it every time we send a message if file_type == ChatFileType.DOC: extracted_text = extract_file_text( - file=file_content_io, # use the bytes we already read + file=file_content_io, # use the bytes we already read file_name=file.filename or "", ) text_file_id = str(uuid.uuid4()) diff --git a/backend/onyx/tools/tool_implementations/search/search_tool.py b/backend/onyx/tools/tool_implementations/search/search_tool.py index 2e541f5126b..2666b2014a4 100644 --- a/backend/onyx/tools/tool_implementations/search/search_tool.py +++ b/backend/onyx/tools/tool_implementations/search/search_tool.py @@ -7,7 +7,6 @@ from sqlalchemy.orm import Session from onyx.chat.chat_utils import llm_doc_from_inference_section -from onyx.chat.llm_response_handler import LLMCall from onyx.chat.models import AnswerStyleConfig from onyx.chat.models import ContextualPruningConfig from onyx.chat.models import DocumentPruningConfig @@ -371,41 +370,6 @@ def build_next_prompt( prompt_config=self.prompt_config, ) - """Other utility functions""" - - @classmethod - def get_search_result( - cls, llm_call: LLMCall - ) -> tuple[list[LlmDoc], list[LlmDoc]] | None: - """ - Returns the final search results and a map of docs to their original search rank (which is what is displayed to user) - """ - if not llm_call.tool_call_info: - return None - - final_search_results = [] - initial_search_results = [] - - for yield_item in llm_call.tool_call_info: - if ( - isinstance(yield_item, ToolResponse) - and yield_item.id == FINAL_CONTEXT_DOCUMENTS_ID - ): - final_search_results = cast(list[LlmDoc], yield_item.response) - elif ( - isinstance(yield_item, ToolResponse) - and yield_item.id == SEARCH_DOC_CONTENT_ID - ): - search_contexts = yield_item.response.contexts - # original_doc_search_rank = 1 - for doc in search_contexts: - if doc.document_id not in initial_search_results: - initial_search_results.append(doc) - - initial_search_results = cast(list[LlmDoc], initial_search_results) - - return final_search_results, initial_search_results - # Allows yielding the same responses as a SearchTool without being a SearchTool. # SearchTool passed in to allow for access to SearchTool properties. diff --git a/backend/tests/integration/tests/dev_apis/test_simple_chat_api.py b/backend/tests/integration/tests/dev_apis/test_simple_chat_api.py index 564c1132d6e..cadb848a97d 100644 --- a/backend/tests/integration/tests/dev_apis/test_simple_chat_api.py +++ b/backend/tests/integration/tests/dev_apis/test_simple_chat_api.py @@ -1,6 +1,5 @@ import json -import pytest import requests from onyx.configs.constants import MessageType @@ -66,9 +65,6 @@ def test_send_message_simple_with_history(reset: None) -> None: assert found_doc["metadata"]["document_id"] == doc.id -@pytest.mark.xfail( - reason="agent search broke this", -) def test_using_reference_docs_with_simple_with_history_api_flow(reset: None) -> None: # Creating an admin user (first user created is automatically an admin) admin_user: DATestUser = UserManager.create(name="admin_user")