Skip to content

fix(chat): post llm loop callback #7309

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Danelegend merged 8 commits into from
Jan 9, 2026
Merged

fix(chat): post llm loop callback #7309

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
020af02
nit
Jan 9, 2026
fde21ed
nits
Jan 9, 2026
1d72eff
replace log with emit
Jan 9, 2026
d273e63
update last turn index
Jan 9, 2026
fa218ac
Fix nit
Jan 9, 2026
cee60dd
Resolve merge conflict
Jan 9, 2026
448e97b
nits
Jan 9, 2026
05b56df
nits
Jan 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions backend/onyx/chat/chat_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ def get_is_clarification(self) -> bool:

def run_chat_loop_with_state_containers(
func: Callable[..., None],
completion_callback: Callable[[ChatStateContainer], None],
is_connected: Callable[[], bool],
emitter: Emitter,
state_container: ChatStateContainer,
Expand Down Expand Up @@ -196,3 +197,4 @@ def run_with_exception_capture() -> None:
# Skip waiting if user disconnected to exit quickly.
if is_connected():
wait_on_background(thread)
completion_callback(state_container)
112 changes: 67 additions & 45 deletions backend/onyx/chat/process_message.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import re
import traceback
from collections.abc import Callable
from uuid import UUID

from sqlalchemy.orm import Session
Expand Down Expand Up @@ -45,6 +46,7 @@
from onyx.db.chat import get_or_create_root_message
from onyx.db.chat import reserve_message_id
from onyx.db.memory import get_memories
from onyx.db.models import ChatMessage
from onyx.db.models import User
from onyx.db.projects import get_project_token_count
from onyx.db.projects import get_user_files_from_project
Expand Down Expand Up @@ -540,6 +542,17 @@ def check_is_connected() -> bool:
# External container allows non-streaming callers to access accumulated state
state_container = external_state_container or ChatStateContainer()

def llm_loop_completion_callback(
state_container: ChatStateContainer,
) -> None:
llm_loop_completion_handle(
state_container=state_container,
db_session=db_session,
chat_session_id=str(chat_session.id),
is_connected=check_is_connected,
assistant_message=assistant_response,
)

# Run the LLM loop with explicit wrapper for stop signal handling
# The wrapper runs run_llm_loop in a background thread and polls every 300ms
# for stop signals. run_llm_loop itself doesn't know about stopping.
Expand All @@ -555,6 +568,7 @@ def check_is_connected() -> bool:

yield from run_chat_loop_with_state_containers(
run_deep_research_llm_loop,
llm_loop_completion_callback,
is_connected=check_is_connected,
emitter=emitter,
state_container=state_container,
Expand All @@ -571,6 +585,7 @@ def check_is_connected() -> bool:
else:
yield from run_chat_loop_with_state_containers(
run_llm_loop,
llm_loop_completion_callback,
is_connected=check_is_connected, # Not passed through to run_llm_loop
emitter=emitter,
state_container=state_container,
Expand All @@ -588,51 +603,6 @@ def check_is_connected() -> bool:
chat_session_id=str(chat_session.id),
)

# Determine if stopped by user
completed_normally = check_is_connected()
if not completed_normally:
logger.debug(f"Chat session {chat_session.id} stopped by user")

# Build final answer based on completion status
if completed_normally:
if state_container.answer_tokens is None:
raise RuntimeError(
"LLM run completed normally but did not return an answer."
)
final_answer = state_container.answer_tokens
else:
# Stopped by user - append stop message
if state_container.answer_tokens:
final_answer = (
state_container.answer_tokens
+ " ... The generation was stopped by the user here."
)
else:
final_answer = "The generation was stopped by the user."

# Build citation_docs_info from accumulated citations in state container
citation_docs_info: list[CitationDocInfo] = []
seen_citation_nums: set[int] = set()
for citation_num, search_doc in state_container.citation_to_doc.items():
if citation_num not in seen_citation_nums:
seen_citation_nums.add(citation_num)
citation_docs_info.append(
CitationDocInfo(
search_doc=search_doc,
citation_number=citation_num,
)
)

save_chat_turn(
message_text=final_answer,
reasoning_tokens=state_container.reasoning_tokens,
citation_docs_info=citation_docs_info,
tool_calls=state_container.tool_calls,
db_session=db_session,
assistant_message=assistant_response,
is_clarification=state_container.is_clarification,
)

except ValueError as e:
logger.exception("Failed to process chat message.")

Expand Down Expand Up @@ -693,6 +663,58 @@ def check_is_connected() -> bool:
return


def llm_loop_completion_handle(
state_container: ChatStateContainer,
is_connected: Callable[[], bool],
db_session: Session,
chat_session_id: str,
assistant_message: ChatMessage,
) -> None:
completed_normally = is_connected()
if not completed_normally:
logger.debug(f"Chat session {chat_session_id} stopped by user")

# Build final answer based on completion status
if completed_normally:
if state_container.answer_tokens is None:
raise RuntimeError(
"LLM run completed normally but did not return an answer."
)
Comment on lines +660 to +664
Copy link
Contributor

@greptile-apps greptile-apps bot Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] When the LLM completes normally but answer_tokens is None, this raises RuntimeError. However, since this runs in the completion callback (which executes in the finally block), this exception will propagate differently than the original design where it was caught by the exception handlers in handle_stream_message_objects. The error handling path has changed.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: backend/onyx/chat/process_message.py
Line: 678:682

Comment:
[P2] When the LLM completes normally but `answer_tokens` is None, this raises RuntimeError. However, since this runs in the completion callback (which executes in the `finally` block), this exception will propagate differently than the original design where it was caught by the exception handlers in `handle_stream_message_objects`. The error handling path has changed.

How can I resolve this? If you propose a fix, please make it concise.

Copy link
Contributor Author

@Danelegend Danelegend Jan 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is handled in the layer that runs this func

final_answer = state_container.answer_tokens
else:
# Stopped by user - append stop message
if state_container.answer_tokens:
final_answer = (
state_container.answer_tokens
+ " ... The generation was stopped by the user here."
)
else:
final_answer = "The generation was stopped by the user."

# Build citation_docs_info from accumulated citations in state container
citation_docs_info: list[CitationDocInfo] = []
seen_citation_nums: set[int] = set()
for citation_num, search_doc in state_container.citation_to_doc.items():
if citation_num not in seen_citation_nums:
seen_citation_nums.add(citation_num)
citation_docs_info.append(
CitationDocInfo(
search_doc=search_doc,
citation_number=citation_num,
)
)

save_chat_turn(
message_text=final_answer,
reasoning_tokens=state_container.reasoning_tokens,
citation_docs_info=citation_docs_info,
tool_calls=state_container.tool_calls,
db_session=db_session,
assistant_message=assistant_message,
is_clarification=state_container.is_clarification,
)


def stream_chat_message_objects(
new_msg_req: CreateChatMessageRequest,
user: User | None,
Expand Down
Loading