diff --git a/rag/app/qa.py b/rag/app/qa.py index f663d7142e..156ddda9f6 100644 --- a/rag/app/qa.py +++ b/rag/app/qa.py @@ -353,12 +353,16 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang= fails = [] question, answer = "", "" res = [] - reader = csv.reader(lines, delimiter=delimiter) + reader = csv.reader((line + "\n" for line in lines), delimiter=delimiter) + prev_line_num = 0 + # line_num tracks the physical span when quoted fields cross lines. for i, row in enumerate(reader): + raw = "\n".join(lines[prev_line_num : reader.line_num]) + prev_line_num = reader.line_num if len(row) != 2: if question: - answer += "\n" + lines[i] + answer += "\n" + raw else: fails.append(str(i + 1)) elif len(row) == 2: diff --git a/test/unit_test/rag/app/test_qa_csv.py b/test/unit_test/rag/app/test_qa_csv.py index 04897fe65b..b23ec81bf7 100644 --- a/test/unit_test/rag/app/test_qa_csv.py +++ b/test/unit_test/rag/app/test_qa_csv.py @@ -52,3 +52,42 @@ def test_csv_final_pair_uses_last_line_number(): assert len(chunks) == 2 assert chunks[0]["top_int"] == [1] assert chunks[1]["top_int"] == [2] + + +@pytest.mark.p2 +def test_csv_quoted_comma_stays_in_one_field(): + chunks = qa.chunk( + "qa.csv", + binary=b'"Question, one",Answer 1\nQuestion 2,Answer 2', + lang="English", + callback=_noop_callback, + ) + + assert len(chunks) == 2 + assert chunks[0]["content_with_weight"] == "Question: Question, one\tAnswer: 1" + + +@pytest.mark.p2 +def test_csv_quoted_field_preserves_embedded_newline(): + chunks = qa.chunk( + "qa.csv", + binary=b'"first line\nsecond line",some answer\n', + lang="English", + callback=_noop_callback, + ) + + assert len(chunks) == 1 + assert chunks[0]["content_with_weight"] == "Question: first line\nsecond line\tAnswer: some answer" + + +@pytest.mark.p2 +def test_csv_multiline_quote_uses_physical_continuation_line(): + chunks = qa.chunk( + "qa.csv", + binary=b'first,one\n"second\nquestion",two\ntwo continued\nthird,three', + lang="English", + callback=_noop_callback, + ) + + assert len(chunks) == 3 + assert chunks[1]["content_with_weight"] == "Question: second\nquestion\tAnswer: two\ntwo continued"