mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 19:32:55 -05:00
459 lines
18 KiB
Python
459 lines
18 KiB
Python
from collections.abc import Callable
|
|
from collections.abc import Iterator
|
|
from functools import partial
|
|
from typing import cast
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from danswer.chat.chat_utils import create_chat_chain
|
|
from danswer.chat.chat_utils import llm_doc_from_inference_section
|
|
from danswer.chat.models import CitationInfo
|
|
from danswer.chat.models import DanswerAnswerPiece
|
|
from danswer.chat.models import LlmDoc
|
|
from danswer.chat.models import LLMRelevanceFilterResponse
|
|
from danswer.chat.models import QADocsResponse
|
|
from danswer.chat.models import StreamingError
|
|
from danswer.configs.chat_configs import CHAT_TARGET_CHUNK_PERCENTAGE
|
|
from danswer.configs.chat_configs import MAX_CHUNKS_FED_TO_CHAT
|
|
from danswer.configs.constants import MessageType
|
|
from danswer.db.chat import create_db_search_doc
|
|
from danswer.db.chat import create_new_chat_message
|
|
from danswer.db.chat import get_chat_message
|
|
from danswer.db.chat import get_chat_session_by_id
|
|
from danswer.db.chat import get_db_search_doc_by_id
|
|
from danswer.db.chat import get_doc_query_identifiers_from_model
|
|
from danswer.db.chat import get_or_create_root_message
|
|
from danswer.db.chat import translate_db_message_to_chat_message_detail
|
|
from danswer.db.chat import translate_db_search_doc_to_server_search_doc
|
|
from danswer.db.embedding_model import get_current_db_embedding_model
|
|
from danswer.db.engine import get_session_context_manager
|
|
from danswer.db.models import SearchDoc as DbSearchDoc
|
|
from danswer.db.models import User
|
|
from danswer.document_index.factory import get_default_document_index
|
|
from danswer.file_store.models import ChatFileType
|
|
from danswer.file_store.utils import load_all_chat_files
|
|
from danswer.llm.answering.answer import Answer
|
|
from danswer.llm.answering.models import AnswerStyleConfig
|
|
from danswer.llm.answering.models import CitationConfig
|
|
from danswer.llm.answering.models import DocumentPruningConfig
|
|
from danswer.llm.answering.models import PreviousMessage
|
|
from danswer.llm.answering.models import PromptConfig
|
|
from danswer.llm.exceptions import GenAIDisabledException
|
|
from danswer.llm.factory import get_llm_for_persona
|
|
from danswer.llm.utils import get_default_llm_tokenizer
|
|
from danswer.search.models import OptionalSearchSetting
|
|
from danswer.search.models import SearchRequest
|
|
from danswer.search.pipeline import SearchPipeline
|
|
from danswer.search.retrieval.search_runner import inference_documents_from_ids
|
|
from danswer.search.utils import chunks_or_sections_to_search_docs
|
|
from danswer.secondary_llm_flows.choose_search import check_if_need_search
|
|
from danswer.secondary_llm_flows.query_expansion import history_based_query_rephrase
|
|
from danswer.server.query_and_chat.models import ChatMessageDetail
|
|
from danswer.server.query_and_chat.models import CreateChatMessageRequest
|
|
from danswer.server.utils import get_json_line
|
|
from danswer.utils.logger import setup_logger
|
|
from danswer.utils.timing import log_generator_function_time
|
|
|
|
logger = setup_logger()
|
|
|
|
|
|
def translate_citations(
|
|
citations_list: list[CitationInfo], db_docs: list[DbSearchDoc]
|
|
) -> dict[int, int]:
|
|
"""Always cites the first instance of the document_id, assumes the db_docs
|
|
are sorted in the order displayed in the UI"""
|
|
doc_id_to_saved_doc_id_map: dict[str, int] = {}
|
|
for db_doc in db_docs:
|
|
if db_doc.document_id not in doc_id_to_saved_doc_id_map:
|
|
doc_id_to_saved_doc_id_map[db_doc.document_id] = db_doc.id
|
|
|
|
citation_to_saved_doc_id_map: dict[int, int] = {}
|
|
for citation in citations_list:
|
|
if citation.citation_num not in citation_to_saved_doc_id_map:
|
|
citation_to_saved_doc_id_map[
|
|
citation.citation_num
|
|
] = doc_id_to_saved_doc_id_map[citation.document_id]
|
|
|
|
return citation_to_saved_doc_id_map
|
|
|
|
|
|
ChatPacketStream = Iterator[
|
|
StreamingError
|
|
| QADocsResponse
|
|
| LLMRelevanceFilterResponse
|
|
| ChatMessageDetail
|
|
| DanswerAnswerPiece
|
|
| CitationInfo
|
|
]
|
|
|
|
|
|
def stream_chat_message_objects(
|
|
new_msg_req: CreateChatMessageRequest,
|
|
user: User | None,
|
|
db_session: Session,
|
|
# Needed to translate persona num_chunks to tokens to the LLM
|
|
default_num_chunks: float = MAX_CHUNKS_FED_TO_CHAT,
|
|
# For flow with search, don't include as many chunks as possible since we need to leave space
|
|
# for the chat history, for smaller models, we likely won't get MAX_CHUNKS_FED_TO_CHAT chunks
|
|
max_document_percentage: float = CHAT_TARGET_CHUNK_PERCENTAGE,
|
|
# if specified, uses the last user message and does not create a new user message based
|
|
# on the `new_msg_req.message`. Currently, requires a state where the last message is a
|
|
# user message (e.g. this can only be used for the chat-seeding flow).
|
|
use_existing_user_message: bool = False,
|
|
) -> ChatPacketStream:
|
|
"""Streams in order:
|
|
1. [conditional] Retrieved documents if a search needs to be run
|
|
2. [conditional] LLM selected chunk indices if LLM chunk filtering is turned on
|
|
3. [always] A set of streamed LLM tokens or an error anywhere along the line if something fails
|
|
4. [always] Details on the final AI response message that is created
|
|
|
|
"""
|
|
try:
|
|
user_id = user.id if user is not None else None
|
|
|
|
chat_session = get_chat_session_by_id(
|
|
chat_session_id=new_msg_req.chat_session_id,
|
|
user_id=user_id,
|
|
db_session=db_session,
|
|
)
|
|
|
|
message_text = new_msg_req.message
|
|
chat_session_id = new_msg_req.chat_session_id
|
|
parent_id = new_msg_req.parent_message_id
|
|
reference_doc_ids = new_msg_req.search_doc_ids
|
|
retrieval_options = new_msg_req.retrieval_options
|
|
persona = chat_session.persona
|
|
query_override = new_msg_req.query_override
|
|
|
|
# After this section, no_ai_answer is represented by prompt being None
|
|
prompt_id = new_msg_req.prompt_id
|
|
if prompt_id is None and persona.prompts and not new_msg_req.no_ai_answer:
|
|
prompt_id = sorted(persona.prompts, key=lambda x: x.id)[-1].id
|
|
|
|
if reference_doc_ids is None and retrieval_options is None:
|
|
raise RuntimeError(
|
|
"Must specify a set of documents for chat or specify search options"
|
|
)
|
|
|
|
try:
|
|
llm = get_llm_for_persona(
|
|
persona, new_msg_req.llm_override or chat_session.llm_override
|
|
)
|
|
except GenAIDisabledException:
|
|
llm = None
|
|
|
|
llm_tokenizer = get_default_llm_tokenizer()
|
|
llm_tokenizer_encode_func = cast(
|
|
Callable[[str], list[int]], llm_tokenizer.encode
|
|
)
|
|
|
|
embedding_model = get_current_db_embedding_model(db_session)
|
|
document_index = get_default_document_index(
|
|
primary_index_name=embedding_model.index_name, secondary_index_name=None
|
|
)
|
|
|
|
# Every chat Session begins with an empty root message
|
|
root_message = get_or_create_root_message(
|
|
chat_session_id=chat_session_id, db_session=db_session
|
|
)
|
|
|
|
if parent_id is not None:
|
|
parent_message = get_chat_message(
|
|
chat_message_id=parent_id,
|
|
user_id=user_id,
|
|
db_session=db_session,
|
|
)
|
|
else:
|
|
parent_message = root_message
|
|
|
|
if not use_existing_user_message:
|
|
# Create new message at the right place in the tree and update the parent's child pointer
|
|
# Don't commit yet until we verify the chat message chain
|
|
user_message = create_new_chat_message(
|
|
chat_session_id=chat_session_id,
|
|
parent_message=parent_message,
|
|
prompt_id=prompt_id,
|
|
message=message_text,
|
|
token_count=len(llm_tokenizer_encode_func(message_text)),
|
|
message_type=MessageType.USER,
|
|
files=[
|
|
{"id": str(file_id), "type": ChatFileType.IMAGE}
|
|
for file_id in new_msg_req.file_ids
|
|
],
|
|
db_session=db_session,
|
|
commit=False,
|
|
)
|
|
# re-create linear history of messages
|
|
final_msg, history_msgs = create_chat_chain(
|
|
chat_session_id=chat_session_id, db_session=db_session
|
|
)
|
|
if final_msg.id != user_message.id:
|
|
db_session.rollback()
|
|
raise RuntimeError(
|
|
"The new message was not on the mainline. "
|
|
"Be sure to update the chat pointers before calling this."
|
|
)
|
|
|
|
# Save now to save the latest chat message
|
|
db_session.commit()
|
|
else:
|
|
# re-create linear history of messages
|
|
final_msg, history_msgs = create_chat_chain(
|
|
chat_session_id=chat_session_id, db_session=db_session
|
|
)
|
|
if final_msg.message_type != MessageType.USER:
|
|
raise RuntimeError(
|
|
"The last message was not a user message. Cannot call "
|
|
"`stream_chat_message_objects` with `is_regenerate=True` "
|
|
"when the last message is not a user message."
|
|
)
|
|
|
|
# load all files needed for this chat chain in memory
|
|
files = load_all_chat_files(history_msgs, new_msg_req.file_ids, db_session)
|
|
latest_query_files = [
|
|
file for file in files if file.file_id in new_msg_req.file_ids
|
|
]
|
|
|
|
run_search = False
|
|
# Retrieval options are only None if reference_doc_ids are provided
|
|
# Also don't perform search if the user uploaded at least one file - just use the files
|
|
if (
|
|
retrieval_options is not None
|
|
and persona.num_chunks != 0
|
|
and not new_msg_req.file_ids
|
|
):
|
|
if retrieval_options.run_search == OptionalSearchSetting.ALWAYS:
|
|
run_search = True
|
|
elif retrieval_options.run_search == OptionalSearchSetting.NEVER:
|
|
run_search = False
|
|
else:
|
|
run_search = check_if_need_search(
|
|
query_message=final_msg, history=history_msgs, llm=llm
|
|
)
|
|
|
|
rephrased_query = None
|
|
llm_relevance_list = None
|
|
if reference_doc_ids:
|
|
identifier_tuples = get_doc_query_identifiers_from_model(
|
|
search_doc_ids=reference_doc_ids,
|
|
chat_session=chat_session,
|
|
user_id=user_id,
|
|
db_session=db_session,
|
|
)
|
|
|
|
# Generates full documents currently
|
|
# May extend to include chunk ranges
|
|
llm_docs: list[LlmDoc] = inference_documents_from_ids(
|
|
doc_identifiers=identifier_tuples,
|
|
document_index=document_index,
|
|
)
|
|
document_pruning_config = DocumentPruningConfig(
|
|
is_manually_selected_docs=True
|
|
)
|
|
|
|
# In case the search doc is deleted, just don't include it
|
|
# though this should never happen
|
|
db_search_docs_or_none = [
|
|
get_db_search_doc_by_id(doc_id=doc_id, db_session=db_session)
|
|
for doc_id in reference_doc_ids
|
|
]
|
|
|
|
reference_db_search_docs = [
|
|
db_sd for db_sd in db_search_docs_or_none if db_sd
|
|
]
|
|
|
|
elif run_search:
|
|
rephrased_query = (
|
|
history_based_query_rephrase(
|
|
query_message=final_msg, history=history_msgs, llm=llm
|
|
)
|
|
if query_override is None
|
|
else query_override
|
|
)
|
|
|
|
search_pipeline = SearchPipeline(
|
|
search_request=SearchRequest(
|
|
query=rephrased_query,
|
|
human_selected_filters=retrieval_options.filters
|
|
if retrieval_options
|
|
else None,
|
|
persona=persona,
|
|
offset=retrieval_options.offset if retrieval_options else None,
|
|
limit=retrieval_options.limit if retrieval_options else None,
|
|
chunks_above=new_msg_req.chunks_above,
|
|
chunks_below=new_msg_req.chunks_below,
|
|
full_doc=new_msg_req.full_doc,
|
|
),
|
|
user=user,
|
|
db_session=db_session,
|
|
)
|
|
|
|
top_sections = search_pipeline.reranked_sections
|
|
top_docs = chunks_or_sections_to_search_docs(top_sections)
|
|
|
|
reference_db_search_docs = [
|
|
create_db_search_doc(server_search_doc=top_doc, db_session=db_session)
|
|
for top_doc in top_docs
|
|
]
|
|
|
|
response_docs = [
|
|
translate_db_search_doc_to_server_search_doc(db_search_doc)
|
|
for db_search_doc in reference_db_search_docs
|
|
]
|
|
|
|
initial_response = QADocsResponse(
|
|
rephrased_query=rephrased_query,
|
|
top_documents=response_docs,
|
|
predicted_flow=search_pipeline.predicted_flow,
|
|
predicted_search=search_pipeline.predicted_search_type,
|
|
applied_source_filters=search_pipeline.search_query.filters.source_type,
|
|
applied_time_cutoff=search_pipeline.search_query.filters.time_cutoff,
|
|
recency_bias_multiplier=search_pipeline.search_query.recency_bias_multiplier,
|
|
)
|
|
yield initial_response
|
|
|
|
# Yield the list of LLM selected chunks for showing the LLM selected icons in the UI
|
|
llm_relevance_filtering_response = LLMRelevanceFilterResponse(
|
|
relevant_chunk_indices=search_pipeline.relevant_chunk_indices
|
|
)
|
|
yield llm_relevance_filtering_response
|
|
|
|
document_pruning_config = DocumentPruningConfig(
|
|
max_chunks=int(
|
|
persona.num_chunks
|
|
if persona.num_chunks is not None
|
|
else default_num_chunks
|
|
),
|
|
max_window_percentage=max_document_percentage,
|
|
use_sections=search_pipeline.ran_merge_chunk,
|
|
)
|
|
|
|
llm_docs = [
|
|
llm_doc_from_inference_section(section) for section in top_sections
|
|
]
|
|
llm_relevance_list = search_pipeline.section_relevance_list
|
|
|
|
else:
|
|
llm_docs = []
|
|
reference_db_search_docs = None
|
|
document_pruning_config = DocumentPruningConfig()
|
|
|
|
# Cannot determine these without the LLM step or breaking out early
|
|
partial_response = partial(
|
|
create_new_chat_message,
|
|
chat_session_id=chat_session_id,
|
|
parent_message=final_msg,
|
|
prompt_id=prompt_id,
|
|
# message=,
|
|
rephrased_query=rephrased_query,
|
|
# token_count=,
|
|
message_type=MessageType.ASSISTANT,
|
|
# error=,
|
|
reference_docs=reference_db_search_docs,
|
|
db_session=db_session,
|
|
commit=True,
|
|
)
|
|
|
|
# If no prompt is provided, this is interpreted as not wanting an AI Answer
|
|
# Simply provide/save the retrieval results
|
|
if final_msg.prompt is None:
|
|
gen_ai_response_message = partial_response(
|
|
message="",
|
|
token_count=0,
|
|
citations=None,
|
|
error=None,
|
|
)
|
|
msg_detail_response = translate_db_message_to_chat_message_detail(
|
|
gen_ai_response_message
|
|
)
|
|
|
|
yield msg_detail_response
|
|
|
|
# Stop here after saving message details, the above still needs to be sent for the
|
|
# message id to send the next follow-up message
|
|
return
|
|
|
|
# LLM prompt building, response capturing, etc.
|
|
answer = Answer(
|
|
question=final_msg.message,
|
|
docs=llm_docs,
|
|
latest_query_files=latest_query_files,
|
|
answer_style_config=AnswerStyleConfig(
|
|
citation_config=CitationConfig(
|
|
all_docs_useful=reference_db_search_docs is not None
|
|
),
|
|
document_pruning_config=document_pruning_config,
|
|
),
|
|
prompt_config=PromptConfig.from_model(
|
|
final_msg.prompt,
|
|
prompt_override=(
|
|
new_msg_req.prompt_override or chat_session.prompt_override
|
|
),
|
|
),
|
|
llm=(
|
|
llm
|
|
or get_llm_for_persona(
|
|
persona, new_msg_req.llm_override or chat_session.llm_override
|
|
)
|
|
),
|
|
doc_relevance_list=llm_relevance_list,
|
|
message_history=[
|
|
PreviousMessage.from_chat_message(msg, files) for msg in history_msgs
|
|
],
|
|
)
|
|
# generator will not include quotes, so we can cast
|
|
yield from cast(ChatPacketStream, answer.processed_streamed_output)
|
|
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
|
|
# Frontend will erase whatever answer and show this instead
|
|
# This will be the issue 99% of the time
|
|
yield StreamingError(error="LLM failed to respond, have you set your API key?")
|
|
return
|
|
|
|
# Post-LLM answer processing
|
|
try:
|
|
db_citations = None
|
|
if reference_db_search_docs:
|
|
db_citations = translate_citations(
|
|
citations_list=answer.citations,
|
|
db_docs=reference_db_search_docs,
|
|
)
|
|
|
|
# Saving Gen AI answer and responding with message info
|
|
gen_ai_response_message = partial_response(
|
|
message=answer.llm_answer,
|
|
token_count=len(llm_tokenizer_encode_func(answer.llm_answer)),
|
|
citations=db_citations,
|
|
error=None,
|
|
)
|
|
|
|
msg_detail_response = translate_db_message_to_chat_message_detail(
|
|
gen_ai_response_message
|
|
)
|
|
|
|
yield msg_detail_response
|
|
except Exception as e:
|
|
logger.exception(e)
|
|
|
|
# Frontend will erase whatever answer and show this instead
|
|
yield StreamingError(error="Failed to parse LLM output")
|
|
|
|
|
|
@log_generator_function_time()
|
|
def stream_chat_message(
|
|
new_msg_req: CreateChatMessageRequest,
|
|
user: User | None,
|
|
use_existing_user_message: bool = False,
|
|
) -> Iterator[str]:
|
|
with get_session_context_manager() as db_session:
|
|
objects = stream_chat_message_objects(
|
|
new_msg_req=new_msg_req,
|
|
user=user,
|
|
db_session=db_session,
|
|
use_existing_user_message=use_existing_user_message,
|
|
)
|
|
for obj in objects:
|
|
yield get_json_line(obj.dict())
|