add unit test for quotes (#2085)

* add unit test for quotes

* test answer and quotes together

backend/tests/unit/danswer/llm/answering/stream_processing/test_quote_processing.py

@@ -2,6 +2,7 @@ import json
 from datetime import datetime
 
 from danswer.chat.models import DanswerAnswerPiece
+from danswer.chat.models import DanswerQuotes
 from danswer.chat.models import LlmDoc
 from danswer.configs.constants import DocumentSource
 from danswer.llm.answering.stream_processing.quotes_processing import (
@@ -24,8 +25,7 @@ mock_docs = [
 ]
 
 
-def test_process_model_tokens() -> None:
+tokens_with_quotes = [
-    tokens = [
     "{",
     "\n  ",
     '"answer": "Yes',
@@ -182,10 +182,13 @@ def test_process_model_tokens() -> None:
     ' way."\n  ]',
     "\n}",
     "",
-    ]
+]
-    gen = process_model_tokens(tokens=iter(tokens), context_docs=mock_docs)
 
-    s_json = "".join(tokens)
+
+def test_process_model_tokens_answer() -> None:
+    gen = process_model_tokens(tokens=iter(tokens_with_quotes), context_docs=mock_docs)
+
+    s_json = "".join(tokens_with_quotes)
     j = json.loads(s_json)
     expected_answer = j["answer"]
     actual = ""
@@ -278,3 +281,37 @@ def test_json_answer_split_tokens() -> None:
     )
 
     assert expected_answer == actual
+
+
+def test_prefixed_json_with_quotes() -> None:
+    tokens = [
+        "```",
+        "json",
+        "\n",
+        "{",
+        '"answer": "This is a simple ',
+        "answer.",
+        '",\n"',
+        'quotes": ["Document"]',
+        "\n}",
+        "\n",
+        "```",
+    ]
+
+    gen = process_model_tokens(tokens=iter(tokens), context_docs=mock_docs)
+
+    actual_answer = ""
+    actual_count = 0
+    for o in gen:
+        if isinstance(o, DanswerAnswerPiece):
+            if o.answer_piece:
+                actual_answer += o.answer_piece
+            continue
+
+        if isinstance(o, DanswerQuotes):
+            for q in o.quotes:
+                assert q.quote == "Document"
+                actual_count += 1
+
+    assert "This is a simple answer." == actual_answer
+    assert 1 == actual_count