Add back llm_chunks_indices (#2491)

This commit is contained in:
Chris Weaver 2024-09-17 18:21:31 -07:00 committed by GitHub
parent a9427f190a
commit 5f25b243c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 51 additions and 66 deletions

View File

@ -14,6 +14,7 @@ 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.chat.process_message import ChatPacketStream
from danswer.chat.process_message import stream_chat_message_objects
from danswer.configs.constants import MessageType
from danswer.configs.danswerbot_configs import DANSWER_BOT_TARGET_CHUNK_PERCENTAGE
@ -81,6 +82,49 @@ def _get_final_context_doc_indices(
]
def _convert_packet_stream_to_response(
packets: ChatPacketStream,
) -> ChatBasicResponse:
response = ChatBasicResponse()
final_context_docs: list[LlmDoc] = []
answer = ""
for packet in packets:
if isinstance(packet, DanswerAnswerPiece) and packet.answer_piece:
answer += packet.answer_piece
elif isinstance(packet, QADocsResponse):
response.top_documents = packet.top_documents
# TODO: deprecate `simple_search_docs`
response.simple_search_docs = _translate_doc_response_to_simple_doc(packet)
elif isinstance(packet, StreamingError):
response.error_msg = packet.error
elif isinstance(packet, ChatMessageDetail):
response.message_id = packet.message_id
elif isinstance(packet, LLMRelevanceFilterResponse):
response.llm_selected_doc_indices = packet.llm_selected_doc_indices
# TODO: deprecate `llm_chunks_indices`
response.llm_chunks_indices = packet.llm_selected_doc_indices
elif isinstance(packet, FinalUsedContextDocsResponse):
final_context_docs = packet.final_context_docs
elif isinstance(packet, AllCitations):
response.cited_documents = {
citation.citation_num: citation.document_id
for citation in packet.citations
}
response.final_context_doc_indices = _get_final_context_doc_indices(
final_context_docs, response.top_documents
)
response.answer = answer
if answer:
response.answer_citationless = remove_answer_citations(answer)
return response
def remove_answer_citations(answer: str) -> str:
pattern = r"\s*\[\[\d+\]\]\(http[s]?://[^\s]+\)"
@ -140,37 +184,7 @@ def handle_simplified_chat_message(
db_session=db_session,
)
response = ChatBasicResponse()
final_context_docs: list[LlmDoc] = []
answer = ""
for packet in packets:
if isinstance(packet, DanswerAnswerPiece) and packet.answer_piece:
answer += packet.answer_piece
elif isinstance(packet, QADocsResponse):
response.simple_search_docs = _translate_doc_response_to_simple_doc(packet)
response.top_documents = packet.top_documents
elif isinstance(packet, StreamingError):
response.error_msg = packet.error
elif isinstance(packet, ChatMessageDetail):
response.message_id = packet.message_id
elif isinstance(packet, FinalUsedContextDocsResponse):
final_context_docs = packet.final_context_docs
elif isinstance(packet, AllCitations):
response.cited_documents = {
citation.citation_num: citation.document_id
for citation in packet.citations
}
response.final_context_doc_indices = _get_final_context_doc_indices(
final_context_docs, response.top_documents
)
response.answer = answer
if answer:
response.answer_citationless = remove_answer_citations(answer)
return response
return _convert_packet_stream_to_response(packets)
@router.post("/send-message-simple-with-history")
@ -289,36 +303,4 @@ def handle_send_message_simple_with_history(
db_session=db_session,
)
response = ChatBasicResponse()
final_context_docs: list[LlmDoc] = []
answer = ""
for packet in packets:
if isinstance(packet, DanswerAnswerPiece) and packet.answer_piece:
answer += packet.answer_piece
elif isinstance(packet, QADocsResponse):
response.simple_search_docs = _translate_doc_response_to_simple_doc(packet)
response.top_documents = packet.top_documents
elif isinstance(packet, StreamingError):
response.error_msg = packet.error
elif isinstance(packet, ChatMessageDetail):
response.message_id = packet.message_id
elif isinstance(packet, LLMRelevanceFilterResponse):
response.llm_selected_doc_indices = packet.llm_selected_doc_indices
elif isinstance(packet, FinalUsedContextDocsResponse):
final_context_docs = packet.final_context_docs
elif isinstance(packet, AllCitations):
response.cited_documents = {
citation.citation_num: citation.document_id
for citation in packet.citations
}
response.final_context_doc_indices = _get_final_context_doc_indices(
final_context_docs, response.top_documents
)
response.answer = answer
if answer:
response.answer_citationless = remove_answer_citations(answer)
return response
return _convert_packet_stream_to_response(packets)

View File

@ -75,8 +75,6 @@ class ChatBasicResponse(BaseModel):
answer: str | None = None
answer_citationless: str | None = None
# TODO: deprecate `simple_search_docs`
simple_search_docs: list[SimpleDoc] | None = None
top_documents: list[SavedSearchDoc] | None = None
error_msg: str | None = None
@ -85,3 +83,8 @@ class ChatBasicResponse(BaseModel):
final_context_doc_indices: list[int] | None = None
# this is a map of the citation number to the document id
cited_documents: dict[int, str] | None = None
# FOR BACKWARDS COMPATIBILITY
# TODO: deprecate both of these
simple_search_docs: list[SimpleDoc] | None = None
llm_chunks_indices: list[int] | None = None