fix citations

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)

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 <number> (link)
+        - [DOCUMENT <number>] (link)
+        - DOCUMENT <number>
+        - [DOCUMENT <number>]
+
+        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]: