diff --git a/backend/danswer/chat/chat_utils.py b/backend/danswer/chat/chat_utils.py index f4b0b2e02c5..c02e49efd44 100644 --- a/backend/danswer/chat/chat_utils.py +++ b/backend/danswer/chat/chat_utils.py @@ -114,27 +114,29 @@ def combine_message_chain( return "\n\n".join(message_strs) -def reorganize_citations( - answer: str, citations: list[CitationInfo] -) -> tuple[str, list[CitationInfo]]: - """For a complete, citation-aware response, we want to reorganize the citations so that +def reorganize_citations(answer: str, citations: list) -> tuple[str, list]: + """ + For a complete, citation-aware response, we want to reorganize the citations so that they are in the order of the documents that were used in the response. This just looks nicer / avoids - confusion ("Why is there [7] when only 2 documents are cited?").""" + confusion ("Why is there [7] when only 2 documents are cited?"). - # Regular expression to find all instances of [[x]](LINK) - pattern = r"\[\[(.*?)\]\]\((.*?)\)" + Now also handles citations in the format [number] in addition to [[number]](LINK). + """ + + pattern = r"\[\[(\d+)\]\]\((.*?)\)|\[(\d+)\]" all_citation_matches = re.findall(pattern, answer) new_citation_info: dict[int, CitationInfo] = {} for citation_match in all_citation_matches: try: - citation_num = int(citation_match[0]) + citation_str = citation_match[0] if citation_match[0] else citation_match[2] + citation_num = int(citation_str) if citation_num in new_citation_info: continue matching_citation = next( - iter([c for c in citations if c.citation_num == int(citation_num)]), + (c for c in citations if c.citation_num == citation_num), None, ) if matching_citation is None: @@ -146,19 +148,30 @@ def reorganize_citations( ) except Exception: pass - # Function to replace citations with their new number def slack_link_format(match: re.Match) -> str: - link_text = match.group(1) - try: - citation_num = int(link_text) - if citation_num in new_citation_info: - link_text = new_citation_info[citation_num].citation_num - except Exception: - pass - - link_url = match.group(2) - return f"[[{link_text}]]({link_url})" + # Case 1: Linked citation ([[number]](LINK)) + if match.group(1): + link_text = match.group(1) + try: + citation_num = int(link_text) + if citation_num in new_citation_info: + link_text = new_citation_info[citation_num].citation_num + except Exception: + pass + link_url = match.group(2) + return f"[[{link_text}]]({link_url})" + # Case 2: Non-linked citation ([number]) + elif match.group(3): + try: + citation_num = int(match.group(3)) + if citation_num in new_citation_info: + citation_num = new_citation_info[citation_num].citation_num + except Exception: + pass + return f"[{citation_num}]" + else: + return match.group(0) # Substitute all matches in the input text new_answer = re.sub(pattern, slack_link_format, answer) diff --git a/backend/danswer/llm/answering/answer.py b/backend/danswer/llm/answering/answer.py index 6a250d02d2a..694764a6c0a 100644 --- a/backend/danswer/llm/answering/answer.py +++ b/backend/danswer/llm/answering/answer.py @@ -1,4 +1,5 @@ from collections.abc import Iterator +import re from typing import cast from uuid import uuid4 @@ -382,6 +383,47 @@ def _raw_output_for_non_explicit_tool_calling_llms( prompt = prompt_builder.build() yield from message_generator_to_string_generator(self.llm.stream(prompt=prompt)) + def _fix_document_references(self, answer: str) -> str: + """ + Searches the input string for DOCUMENT references in any of these forms: + - DOCUMENT (link) + - [DOCUMENT ] (link) + - DOCUMENT + - [DOCUMENT ] + + and converts them to the proper citation format: + - If a link is provided, returns a linked citation: [[number]](link) + - Otherwise, returns a non-linked citation: [number] + + However, if an adjacent citation (linked or non-linked) for the same number already follows + immediately (ignoring whitespace), the DOCUMENT reference is not converted (i.e. it is removed) + to avoid duplicate citations. + """ + pattern = r"\[?DOCUMENT\s+(\d+)\]?(?:\s*\((.*?)\))?" + + def replacer(match: re.Match) -> str: + try: + num = int(match.group(1)) + except Exception: + return match.group(0) + + if match.group(2) and match.group(2).strip(): + citation = f"[[{num}]]({match.group(2).strip()})" + else: + citation = f"[{num}]" + + post_text = answer[match.end():] + adj_pattern = ( + r"^\s*(\[\[\s*" + re.escape(str(num)) + r"\s*\]\]\([^)]+\)|\[\s*" + re.escape(str(num)) + r"\s*\])" + ) + if re.match(adj_pattern, post_text): + # If an adjacent citation for the same number exists, return an empty string (skip replacement). + return "" + else: + return citation + + return re.sub(pattern, replacer, answer) + @property def processed_streamed_output(self) -> AnswerStream: if self._processed_stream is not None: @@ -465,7 +507,7 @@ def llm_answer(self) -> str: if isinstance(packet, DanswerAnswerPiece) and packet.answer_piece: answer += packet.answer_piece - return answer + return self._fix_document_references(answer) @property def citations(self) -> list[CitationInfo]: