From e304ec4ab6d69d98e5d55045deb115dee04f2df4 Mon Sep 17 00:00:00 2001 From: evan-danswer Date: Wed, 19 Feb 2025 15:52:16 -0800 Subject: [PATCH] Agent search history displayed answer (#4052) --- backend/onyx/chat/chat_utils.py | 3 ++- backend/onyx/chat/models.py | 9 +++++++++ backend/onyx/chat/process_message.py | 13 +++++++++---- backend/onyx/llm/models.py | 2 ++ web/src/app/chat/ChatPage.tsx | 27 +++++++++++++++++++++++++++ web/src/app/chat/interfaces.ts | 9 +++++++++ web/src/app/chat/lib.tsx | 4 +++- 7 files changed, 61 insertions(+), 6 deletions(-) diff --git a/backend/onyx/chat/chat_utils.py b/backend/onyx/chat/chat_utils.py index 4d6618887..81f1f3776 100644 --- a/backend/onyx/chat/chat_utils.py +++ b/backend/onyx/chat/chat_utils.py @@ -190,7 +190,8 @@ def create_chat_chain( and previous_message.message_type == MessageType.ASSISTANT and mainline_messages ): - mainline_messages[-1] = current_message + if current_message.refined_answer_improvement: + mainline_messages[-1] = current_message else: mainline_messages.append(current_message) diff --git a/backend/onyx/chat/models.py b/backend/onyx/chat/models.py index 4ba3d742a..6b1c43437 100644 --- a/backend/onyx/chat/models.py +++ b/backend/onyx/chat/models.py @@ -142,6 +142,15 @@ class MessageResponseIDInfo(BaseModel): reserved_assistant_message_id: int +class AgentMessageIDInfo(BaseModel): + level: int + message_id: int + + +class AgenticMessageResponseIDInfo(BaseModel): + agentic_message_ids: list[AgentMessageIDInfo] + + class StreamingError(BaseModel): error: str stack_trace: str | None = None diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index 65c987f23..a3cce6b7b 100644 --- a/backend/onyx/chat/process_message.py +++ b/backend/onyx/chat/process_message.py @@ -11,6 +11,8 @@ from onyx.agents.agent_search.orchestration.nodes.call_tool import ToolCallExcep from onyx.chat.answer import Answer from onyx.chat.chat_utils import create_chat_chain from onyx.chat.chat_utils import create_temporary_persona +from onyx.chat.models import AgenticMessageResponseIDInfo +from onyx.chat.models import AgentMessageIDInfo from onyx.chat.models import AgentSearchPacket from onyx.chat.models import AllCitations from onyx.chat.models import AnswerPostInfo @@ -308,6 +310,7 @@ ChatPacket = ( | CustomToolResponse | MessageSpecificCitations | MessageResponseIDInfo + | AgenticMessageResponseIDInfo | StreamStopInfo | AgentSearchPacket ) @@ -1035,6 +1038,7 @@ def stream_chat_message_objects( next_level = 1 prev_message = gen_ai_response_message agent_answers = answer.llm_answer_by_level() + agentic_message_ids = [] while next_level in agent_answers: next_answer = agent_answers[next_level] info = info_by_subq[ @@ -1059,17 +1063,18 @@ def stream_chat_message_objects( refined_answer_improvement=refined_answer_improvement, is_agentic=True, ) + agentic_message_ids.append( + AgentMessageIDInfo(level=next_level, message_id=next_answer_message.id) + ) next_level += 1 prev_message = next_answer_message logger.debug("Committing messages") db_session.commit() # actually save user / assistant message - msg_detail_response = translate_db_message_to_chat_message_detail( - gen_ai_response_message - ) + yield AgenticMessageResponseIDInfo(agentic_message_ids=agentic_message_ids) - yield msg_detail_response + yield translate_db_message_to_chat_message_detail(gen_ai_response_message) except Exception as e: error_msg = str(e) logger.exception(error_msg) diff --git a/backend/onyx/llm/models.py b/backend/onyx/llm/models.py index d755a9bd8..925c8bc3f 100644 --- a/backend/onyx/llm/models.py +++ b/backend/onyx/llm/models.py @@ -23,6 +23,7 @@ class PreviousMessage(BaseModel): message_type: MessageType files: list[InMemoryChatFile] tool_call: ToolCallFinalResult | None + refined_answer_improvement: bool | None @classmethod def from_chat_message( @@ -47,6 +48,7 @@ class PreviousMessage(BaseModel): ) if chat_message.tool_call else None, + refined_answer_improvement=chat_message.refined_answer_improvement, ) def to_langchain_msg(self) -> BaseMessage: diff --git a/web/src/app/chat/ChatPage.tsx b/web/src/app/chat/ChatPage.tsx index 9b6823707..d55f33d3d 100644 --- a/web/src/app/chat/ChatPage.tsx +++ b/web/src/app/chat/ChatPage.tsx @@ -23,6 +23,7 @@ import { SubQuestionDetail, constructSubQuestions, DocumentsResponse, + AgenticMessageResponseIDInfo, } from "./interfaces"; import Prism from "prismjs"; @@ -1280,6 +1281,8 @@ export function ChatPage({ let toolCall: ToolCallMetadata | null = null; let isImprovement: boolean | undefined = undefined; let isStreamingQuestions = true; + let includeAgentic = false; + let secondLevelMessageId: number | null = null; let initialFetchDetails: null | { user_message_id: number; @@ -1417,6 +1420,17 @@ export function ChatPage({ resetRegenerationState(); } else { const { user_message_id, frozenMessageMap } = initialFetchDetails; + if (Object.hasOwn(packet, "agentic_message_ids")) { + const agenticMessageIds = (packet as AgenticMessageResponseIDInfo) + .agentic_message_ids; + const level1MessageId = agenticMessageIds.find( + (item) => item.level === 1 + )?.message_id; + if (level1MessageId) { + secondLevelMessageId = level1MessageId; + includeAgentic = true; + } + } setChatState((prevState) => { if (prevState.get(chatSessionIdRef.current!) === "loading") { @@ -1667,6 +1681,19 @@ export function ChatPage({ second_level_generating: second_level_generating, agentic_docs: agenticDocs, }, + ...(includeAgentic + ? [ + { + messageId: secondLevelMessageId!, + message: second_level_answer, + type: "assistant" as const, + files: [], + toolCall: null, + parentMessageId: + initialFetchDetails.assistant_message_id!, + }, + ] + : []), ]); } } diff --git a/web/src/app/chat/interfaces.ts b/web/src/app/chat/interfaces.ts index 343097f8e..aa4970e25 100644 --- a/web/src/app/chat/interfaces.ts +++ b/web/src/app/chat/interfaces.ts @@ -155,6 +155,15 @@ export interface MessageResponseIDInfo { reserved_assistant_message_id: number; } +export interface AgentMessageIDInfo { + level: number; + message_id: number; +} + +export interface AgenticMessageResponseIDInfo { + agentic_message_ids: AgentMessageIDInfo[]; +} + export interface DocumentsResponse { top_documents: OnyxDocument[]; rephrased_query: string | null; diff --git a/web/src/app/chat/lib.tsx b/web/src/app/chat/lib.tsx index bae1b8404..0cd5ac1d5 100644 --- a/web/src/app/chat/lib.tsx +++ b/web/src/app/chat/lib.tsx @@ -25,6 +25,7 @@ import { RetrievalType, StreamingError, ToolCallMetadata, + AgenticMessageResponseIDInfo, } from "./interfaces"; import { Persona } from "../admin/assistants/interfaces"; import { ReadonlyURLSearchParams } from "next/navigation"; @@ -154,7 +155,8 @@ export type PacketType = | AgentAnswerPiece | SubQuestionPiece | ExtendedToolResponse - | RefinedAnswerImprovement; + | RefinedAnswerImprovement + | AgenticMessageResponseIDInfo; export async function* sendMessage({ regenerate,