mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-08-08 14:02:09 +02:00
29f5f4edfa
* removed some dead code and fixed citations when a search request is made with sections selected * fix black formatting issue
65 lines
2.5 KiB
Python
65 lines
2.5 KiB
Python
from collections.abc import Iterator
|
|
from typing import cast
|
|
|
|
from langchain_core.messages import AIMessageChunk
|
|
from langchain_core.messages import BaseMessage
|
|
from langgraph.types import StreamWriter
|
|
|
|
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 (
|
|
PassThroughAnswerResponseHandler,
|
|
)
|
|
from onyx.chat.stream_processing.utils import map_document_id_order
|
|
from onyx.utils.logger import setup_logger
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
def process_llm_stream(
|
|
messages: Iterator[BaseMessage],
|
|
should_stream_answer: bool,
|
|
writer: StreamWriter,
|
|
final_search_results: list[LlmDoc] | None = None,
|
|
displayed_search_results: list[OnyxContext] | list[LlmDoc] | None = None,
|
|
) -> AIMessageChunk:
|
|
tool_call_chunk = AIMessageChunk(content="")
|
|
|
|
if final_search_results and displayed_search_results:
|
|
answer_handler: AnswerResponseHandler = CitationResponseHandler(
|
|
context_docs=final_search_results,
|
|
final_doc_id_to_rank_map=map_document_id_order(final_search_results),
|
|
display_doc_id_to_rank_map=map_document_id_order(displayed_search_results),
|
|
)
|
|
else:
|
|
answer_handler = PassThroughAnswerResponseHandler()
|
|
|
|
full_answer = ""
|
|
# This stream will be the llm answer if no tool is chosen. When a tool is chosen,
|
|
# the stream will contain AIMessageChunks with tool call information.
|
|
for message in messages:
|
|
answer_piece = message.content
|
|
if not isinstance(answer_piece, str):
|
|
# this is only used for logging, so fine to
|
|
# just add the string representation
|
|
answer_piece = str(answer_piece)
|
|
full_answer += answer_piece
|
|
|
|
if isinstance(message, AIMessageChunk) and (
|
|
message.tool_call_chunks or message.tool_calls
|
|
):
|
|
tool_call_chunk += message # type: ignore
|
|
elif should_stream_answer:
|
|
for response_part in answer_handler.handle_response_part(message, []):
|
|
write_custom_event(
|
|
"basic_response",
|
|
response_part,
|
|
writer,
|
|
)
|
|
|
|
logger.debug(f"Full answer: {full_answer}")
|
|
return cast(AIMessageChunk, tool_call_chunk)
|