From 2b2ba5478cd6a8915d2bd674ffe76e0362690a8e Mon Sep 17 00:00:00 2001 From: evan-danswer Date: Mon, 17 Feb 2025 20:20:33 -0800 Subject: [PATCH] new is_agentic flag for chatmessages (#4026) * new is_agentic flag for chatmessages * added cancelled error to db * added cancelled error to returned message --- .../9c00a2bccb83_chat_message_agentic.py | 43 +++++++++++++++++++ backend/onyx/chat/process_message.py | 6 ++- backend/onyx/db/chat.py | 5 ++- backend/onyx/db/models.py | 1 + backend/onyx/server/query_and_chat/models.py | 1 + 5 files changed, 54 insertions(+), 2 deletions(-) create mode 100644 backend/alembic/versions/9c00a2bccb83_chat_message_agentic.py diff --git a/backend/alembic/versions/9c00a2bccb83_chat_message_agentic.py b/backend/alembic/versions/9c00a2bccb83_chat_message_agentic.py new file mode 100644 index 000000000..851f7422d --- /dev/null +++ b/backend/alembic/versions/9c00a2bccb83_chat_message_agentic.py @@ -0,0 +1,43 @@ +"""chat_message_agentic + +Revision ID: 9c00a2bccb83 +Revises: b7a7eee5aa15 +Create Date: 2025-02-17 11:15:43.081150 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "9c00a2bccb83" +down_revision = "b7a7eee5aa15" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # First add the column as nullable + op.add_column("chat_message", sa.Column("is_agentic", sa.Boolean(), nullable=True)) + + # Update existing rows based on presence of SubQuestions + op.execute( + """ + UPDATE chat_message + SET is_agentic = EXISTS ( + SELECT 1 + FROM agent__sub_question + WHERE agent__sub_question.primary_question_id = chat_message.id + ) + WHERE is_agentic IS NULL + """ + ) + + # Make the column non-nullable with a default value of False + op.alter_column( + "chat_message", "is_agentic", nullable=False, server_default=sa.text("false") + ) + + +def downgrade() -> None: + op.drop_column("chat_message", "is_agentic") diff --git a/backend/onyx/chat/process_message.py b/backend/onyx/chat/process_message.py index f04be63d9..15722434d 100644 --- a/backend/onyx/chat/process_message.py +++ b/backend/onyx/chat/process_message.py @@ -146,6 +146,7 @@ from onyx.utils.timing import log_generator_function_time from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR logger = setup_logger() +ERROR_TYPE_CANCELLED = "cancelled" def _translate_citations( @@ -631,6 +632,7 @@ def stream_chat_message_objects( db_session=db_session, commit=False, reserved_message_id=reserved_message_id, + is_agentic=new_msg_req.use_agentic_search, ) prompt_override = new_msg_req.prompt_override or chat_session.prompt_override @@ -1015,7 +1017,7 @@ def stream_chat_message_objects( if info.message_specific_citations else None ), - error=None, + error=ERROR_TYPE_CANCELLED if answer.is_cancelled() else None, tool_call=( ToolCall( tool_id=tool_name_to_tool_id[info.tool_result.tool_name], @@ -1053,7 +1055,9 @@ def stream_chat_message_objects( citations=info.message_specific_citations.citation_map if info.message_specific_citations else None, + error=ERROR_TYPE_CANCELLED if answer.is_cancelled() else None, refined_answer_improvement=refined_answer_improvement, + is_agentic=True, ) next_level += 1 prev_message = next_answer_message diff --git a/backend/onyx/db/chat.py b/backend/onyx/db/chat.py index cfd5bf68b..0da54b8ad 100644 --- a/backend/onyx/db/chat.py +++ b/backend/onyx/db/chat.py @@ -629,6 +629,7 @@ def create_new_chat_message( reserved_message_id: int | None = None, overridden_model: str | None = None, refined_answer_improvement: bool | None = None, + is_agentic: bool = False, ) -> ChatMessage: if reserved_message_id is not None: # Edit existing message @@ -650,7 +651,7 @@ def create_new_chat_message( existing_message.alternate_assistant_id = alternate_assistant_id existing_message.overridden_model = overridden_model existing_message.refined_answer_improvement = refined_answer_improvement - + existing_message.is_agentic = is_agentic new_chat_message = existing_message else: # Create new message @@ -670,6 +671,7 @@ def create_new_chat_message( alternate_assistant_id=alternate_assistant_id, overridden_model=overridden_model, refined_answer_improvement=refined_answer_improvement, + is_agentic=is_agentic, ) db_session.add(new_chat_message) @@ -960,6 +962,7 @@ def translate_db_message_to_chat_message_detail( chat_message.sub_questions ), refined_answer_improvement=chat_message.refined_answer_improvement, + error=chat_message.error, ) return chat_msg_detail diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index b321ae926..d25a5eaec 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -1221,6 +1221,7 @@ class ChatMessage(Base): DateTime(timezone=True), server_default=func.now() ) + is_agentic: Mapped[bool] = mapped_column(Boolean, default=False) refined_answer_improvement: Mapped[bool] = mapped_column(Boolean, nullable=True) chat_session: Mapped[ChatSession] = relationship("ChatSession") diff --git a/backend/onyx/server/query_and_chat/models.py b/backend/onyx/server/query_and_chat/models.py index cdbdf354b..16aa80d5d 100644 --- a/backend/onyx/server/query_and_chat/models.py +++ b/backend/onyx/server/query_and_chat/models.py @@ -240,6 +240,7 @@ class ChatMessageDetail(BaseModel): files: list[FileDescriptor] tool_call: ToolCallFinalResult | None refined_answer_improvement: bool | None = None + error: str | None = None def model_dump(self, *args: list, **kwargs: dict[str, Any]) -> dict[str, Any]: # type: ignore initial_dict = super().model_dump(mode="json", *args, **kwargs) # type: ignore