472 lines
16 KiB
Python
472 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Test the RAG reader implementations."""
|
|
import os
|
|
import json
|
|
from unittest.async_case import IsolatedAsyncioTestCase
|
|
|
|
from agentscope.rag import (
|
|
TextReader,
|
|
PDFReader,
|
|
WordReader,
|
|
ExcelReader,
|
|
PowerPointReader,
|
|
)
|
|
|
|
|
|
class RAGReaderText(IsolatedAsyncioTestCase):
|
|
"""Test cases for RAG reader implementations."""
|
|
|
|
async def test_text_reader(self) -> None:
|
|
"""Test the TextReader implementation."""
|
|
# Split by char
|
|
reader = TextReader(
|
|
chunk_size=10,
|
|
split_by="char",
|
|
)
|
|
docs = await reader(
|
|
text="".join(str(i) for i in range(22)),
|
|
)
|
|
self.assertEqual(len(docs), 4)
|
|
self.assertEqual(
|
|
docs[0].metadata.content["text"],
|
|
"0123456789",
|
|
)
|
|
self.assertEqual(
|
|
docs[1].metadata.content["text"],
|
|
"1011121314",
|
|
)
|
|
self.assertEqual(
|
|
docs[2].metadata.content["text"],
|
|
"1516171819",
|
|
)
|
|
self.assertEqual(
|
|
docs[3].metadata.content["text"],
|
|
"2021",
|
|
)
|
|
|
|
# Split by sentence
|
|
reader = TextReader(
|
|
chunk_size=10,
|
|
split_by="sentence",
|
|
)
|
|
docs = await reader(
|
|
text="012345678910111213. 141516171819! 2021? 22",
|
|
)
|
|
self.assertEqual(
|
|
[_.metadata.content["text"] for _ in docs],
|
|
["0123456789", "10111213.", "1415161718", "19!", "2021?", "22"],
|
|
)
|
|
|
|
docs = await reader(
|
|
text="01234. 56789! 10111213? 14151617..",
|
|
)
|
|
self.assertEqual(
|
|
[_.metadata.content["text"] for _ in docs],
|
|
["01234.", "56789!", "10111213?", "14151617.."],
|
|
)
|
|
|
|
# Split by paragraph
|
|
reader = TextReader(
|
|
chunk_size=5,
|
|
split_by="paragraph",
|
|
)
|
|
docs = await reader(
|
|
text="01234\n\n5678910111213.\n\n\n1415",
|
|
)
|
|
self.assertEqual(
|
|
[_.metadata.content["text"] for _ in docs],
|
|
["01234", "56789", "10111", "213.", "1415"],
|
|
)
|
|
|
|
async def test_pdf_reader(self) -> None:
|
|
"""Test the PDFReader implementation."""
|
|
reader = PDFReader(
|
|
chunk_size=200,
|
|
split_by="sentence",
|
|
)
|
|
pdf_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../examples/functionality/rag/example.pdf",
|
|
)
|
|
docs = await reader(pdf_path=pdf_path)
|
|
self.assertEqual(len(docs), 17)
|
|
self.assertEqual(
|
|
[_.metadata.content["text"] for _ in docs][:2],
|
|
[
|
|
"1\nThe Great Transformations: From Print to Space\n"
|
|
"The invention of the printing press in the 15th century "
|
|
"marked a revolutionary change in \nhuman history.",
|
|
"Johannes Gutenberg's innovation democratized knowledge and "
|
|
"made books \naccessible to the common people.",
|
|
],
|
|
)
|
|
|
|
async def test_word_reader_with_images_and_tables(self) -> None:
|
|
"""Test the WordReader implementation with images and table
|
|
separation."""
|
|
# Test with images and table separation enabled
|
|
reader = WordReader(
|
|
chunk_size=200,
|
|
split_by="sentence",
|
|
include_image=True,
|
|
separate_table=True,
|
|
)
|
|
word_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.docx",
|
|
)
|
|
docs = await reader(word_path=word_path)
|
|
|
|
self.assertListEqual(
|
|
[_.metadata.content["type"] for _ in docs],
|
|
["text"] * 4 + ["image"] * 2 + ["text", "image", "text", "text"],
|
|
)
|
|
|
|
self.assertEqual(
|
|
[_.metadata.content.get("text") for _ in docs],
|
|
[
|
|
"AgentScope\n"
|
|
"标题2\n"
|
|
"This is a test file for AgentScope word reader.",
|
|
"标题3\nTest table:",
|
|
"| Header1 | Header2 | Header3 | Header4 |\n"
|
|
"| --- | --- | --- | --- |\n"
|
|
"| 1 | 2 | 3 | 4 |\n"
|
|
"| 5 | 6 | 7 | 8 |",
|
|
"\nTest list:\nAlice\nBob\nCharlie\nDavid\nTest image:",
|
|
None, # image
|
|
None, # image
|
|
"\nText between images",
|
|
None, # image
|
|
"\nText between image and table",
|
|
"| a | b | c |\n| --- | --- | --- |\n| d\ne | f | g |",
|
|
],
|
|
)
|
|
|
|
self.assertEqual(
|
|
[
|
|
_.metadata.content["source"]["media_type"]
|
|
for _ in docs
|
|
if _.metadata.content["type"] == "image"
|
|
],
|
|
["image/png", "image/png", "image/png"],
|
|
)
|
|
|
|
async def test_excel_reader_with_images_and_tables(self) -> None:
|
|
"""Test the ExcelReader implementation with images and table
|
|
separation."""
|
|
# Test with images and table separation enabled
|
|
reader = ExcelReader(
|
|
chunk_size=200,
|
|
split_by="sentence",
|
|
include_image=True,
|
|
separate_table=True,
|
|
include_cell_coordinates=True,
|
|
table_format="markdown",
|
|
)
|
|
excel_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"test.xlsx",
|
|
)
|
|
docs = await reader(excel_path=excel_path)
|
|
|
|
# Verify document types match expected sequence
|
|
# Expected: table blocks from first sheet, then image (row 9),
|
|
# then table blocks from second sheet
|
|
# Order is based on row positions: table (row 0-5) → image (row 9)
|
|
# → table (row 0-4)
|
|
# Note: with include_cell_coordinates=True, cell coordinates are added
|
|
# to each cell (e.g., [A1], [B1], etc.), which increases text length
|
|
# and results in more chunks
|
|
self.assertListEqual(
|
|
[_.metadata.content["type"] for _ in docs],
|
|
["text"] * 3 + ["image"] * 1 + ["text"] * 5,
|
|
)
|
|
|
|
# Verify exact document content
|
|
doc_texts = [_.metadata.content.get("text") for _ in docs]
|
|
|
|
# Verify sheet headers and table content with cell coordinates
|
|
# First text block should contain Employee Info sheet header and table
|
|
# Note: Due to chunk_size=200, the rows are truncated
|
|
# Order: table (row 0-5) → image (row 9) → table (row 0-4)
|
|
self.assertEqual(
|
|
doc_texts[0],
|
|
"Sheet: Employee Info\n"
|
|
"| [A1] John Smith | [B1] 25 | [C1] Engineering | "
|
|
"[D1] 8000 | [E1] 2020-01-15 |\n"
|
|
"| --- | --- | --- | --- | --- |\n"
|
|
"| [A2] Jane Doe | [B2] 30 | [C2] Sales | "
|
|
"[D2] 12000 | [E2] 2019-03-2",
|
|
)
|
|
# Second text block continues the employee table
|
|
self.assertEqual(
|
|
doc_texts[1],
|
|
"0 |\n"
|
|
"| [A3] Mike \\| Johnson | [B3] 35 | [C3] HR | "
|
|
"[D3] 9000 | [E3] 2021-06-10 |\n"
|
|
"| [A4] Sarah Wilson | [B4] 28 | [C4] Finance | "
|
|
"[D4] 10000 | [E4] 2020-09-05 |\n"
|
|
"| [A5] David Brown | [B5] 32 | [C5] Marketi",
|
|
)
|
|
# Third text block continues the employee table
|
|
self.assertEqual(
|
|
doc_texts[2],
|
|
"ng | [D5] 11000 | [E5] 2018-12-01 |",
|
|
)
|
|
# Image block (text is None)
|
|
self.assertIsNone(doc_texts[3])
|
|
# Fourth text block should contain Product Info sheet header and
|
|
# start of table
|
|
self.assertEqual(
|
|
doc_texts[4],
|
|
"Sheet: Product Info\n"
|
|
"| [A1] Product A | [B1] 100 | [C1] 50 | "
|
|
"[D1] High-quality Product A, suitable for various scenarios.",
|
|
)
|
|
# Remaining blocks continue the product table
|
|
self.assertEqual(
|
|
doc_texts[5],
|
|
"|\n"
|
|
"| --- | --- | --- | --- |\n"
|
|
"| [A2] Product B | [B2] 200 | [C2] 30 | "
|
|
"[D2] Product B offers excellent performance.",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[6],
|
|
"|\n"
|
|
"| [A3] Product C | [B3] 300 | [C3] 20 | "
|
|
"[D3] Product C is a market-leading solution.",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[7],
|
|
"|\n"
|
|
"| [A4] Product D | [B4] 400 | [C4] 40 | "
|
|
"[D4] Product D provides comprehensive functionality.",
|
|
)
|
|
|
|
# Verify image media types
|
|
self.assertEqual(
|
|
[
|
|
_.metadata.content["source"]["media_type"]
|
|
for _ in docs
|
|
if _.metadata.content["type"] == "image"
|
|
],
|
|
["image/png"],
|
|
)
|
|
|
|
async def test_ppt_reader_with_images_and_tables(self) -> None:
|
|
"""Test the PowerPointReader implementation with images and table
|
|
separation."""
|
|
# Test with images and table separation enabled (using defaults)
|
|
reader = PowerPointReader(
|
|
chunk_size=200,
|
|
split_by="sentence",
|
|
separate_table=True,
|
|
)
|
|
ppt_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.pptx",
|
|
)
|
|
docs = await reader(ppt_path=ppt_path)
|
|
|
|
# Verify document types match expected sequence
|
|
# Expected: text blocks from slides, then table, then image
|
|
self.assertListEqual(
|
|
[_.metadata.content["type"] for _ in docs],
|
|
["text"] * 5 + ["image"] * 1 + ["text"],
|
|
)
|
|
|
|
# Verify exact document content
|
|
doc_texts = [_.metadata.content.get("text") for _ in docs]
|
|
|
|
# Verify slide content (with slide tags by default)
|
|
self.assertEqual(
|
|
doc_texts[0],
|
|
"<slide index=1>\nAgentScope\nText content in slide 1\n</slide>",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[1],
|
|
"<slide index=2>\nTitle 2\nText content above table",
|
|
)
|
|
# Table should be extracted as a separate block with Markdown format
|
|
self.assertEqual(
|
|
doc_texts[2],
|
|
"| Name | Age | Career |\n"
|
|
"| --- | --- | --- |\n"
|
|
"| Alice | 25 | Teacher |\n"
|
|
"| Bob | 26 | Doctor |",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[3],
|
|
"Text content below table\n</slide>",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[4],
|
|
"<slide index=3>\nTitle 3\ntext content above image",
|
|
)
|
|
# Image block
|
|
self.assertIsNone(doc_texts[5])
|
|
self.assertEqual(
|
|
doc_texts[6],
|
|
"text content below image\n</slide>",
|
|
)
|
|
|
|
# Verify image media types
|
|
self.assertEqual(
|
|
[
|
|
_.metadata.content["source"]["media_type"]
|
|
for _ in docs
|
|
if _.metadata.content["type"] == "image"
|
|
],
|
|
["image/png"],
|
|
)
|
|
|
|
async def test_ppt_reader_with_json_table_format(self) -> None:
|
|
"""Test the PowerPointReader with JSON table format."""
|
|
|
|
reader = PowerPointReader(
|
|
chunk_size=500,
|
|
split_by="sentence",
|
|
include_image=False,
|
|
separate_table=True,
|
|
table_format="json",
|
|
)
|
|
ppt_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.pptx",
|
|
)
|
|
docs = await reader(ppt_path=ppt_path)
|
|
|
|
# Find the table block
|
|
table_texts = [
|
|
_.metadata.content.get("text")
|
|
for _ in docs
|
|
if _.metadata.content.get("text")
|
|
and "JSON array" in _.metadata.content.get("text", "")
|
|
]
|
|
|
|
# Verify we have a table in JSON format
|
|
self.assertEqual(len(table_texts), 1)
|
|
|
|
# Extract JSON part (after the system-info tag)
|
|
table_text = table_texts[0]
|
|
json_part = table_text.split("\n", 1)[1]
|
|
|
|
# Verify it's valid JSON and matches expected table content
|
|
parsed = json.loads(json_part)
|
|
self.assertEqual(
|
|
parsed,
|
|
[
|
|
["Name", "Age", "Career"],
|
|
["Alice", "25", "Teacher"],
|
|
["Bob", "26", "Doctor"],
|
|
],
|
|
)
|
|
|
|
async def test_ppt_reader_without_image(self) -> None:
|
|
"""Test the PowerPointReader without image extraction."""
|
|
reader = PowerPointReader(
|
|
chunk_size=200,
|
|
split_by="sentence",
|
|
include_image=False,
|
|
separate_table=True,
|
|
)
|
|
ppt_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.pptx",
|
|
)
|
|
docs = await reader(ppt_path=ppt_path)
|
|
|
|
# Verify no image blocks are present
|
|
image_blocks = [
|
|
_ for _ in docs if _.metadata.content["type"] == "image"
|
|
]
|
|
self.assertEqual(len(image_blocks), 0)
|
|
|
|
# All blocks should be text type
|
|
self.assertTrue(
|
|
all(_.metadata.content["type"] == "text" for _ in docs),
|
|
)
|
|
|
|
async def test_ppt_reader_merged_table(self) -> None:
|
|
"""Test the PowerPointReader with merged table
|
|
(separate_table=False)."""
|
|
reader = PowerPointReader(
|
|
chunk_size=500,
|
|
split_by="sentence",
|
|
include_image=False,
|
|
separate_table=False,
|
|
)
|
|
ppt_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.pptx",
|
|
)
|
|
docs = await reader(ppt_path=ppt_path)
|
|
|
|
# When separate_table=False, table should be merged with adjacent text
|
|
# Find the document that contains the table
|
|
table_doc = None
|
|
for doc in docs:
|
|
text = doc.metadata.content.get("text", "")
|
|
if "Name" in text and "Age" in text and "Career" in text:
|
|
table_doc = doc
|
|
break
|
|
|
|
self.assertIsNotNone(table_doc)
|
|
# The table should be merged with surrounding text (with slide tags)
|
|
table_text = table_doc.metadata.content.get("text", "")
|
|
self.assertEqual(
|
|
table_text,
|
|
"<slide index=2>\n"
|
|
"Title 2\n"
|
|
"Text content above table\n"
|
|
"| Name | Age | Career |\n"
|
|
"| --- | --- | --- |\n"
|
|
"| Alice | 25 | Teacher |\n"
|
|
"| Bob | 26 | Doctor |\n"
|
|
"\n"
|
|
"Text content below table\n"
|
|
"</slide>",
|
|
)
|
|
|
|
async def test_ppt_reader_without_slide_tags(self) -> None:
|
|
"""Test the PowerPointReader without slide prefix/suffix XML tags."""
|
|
reader = PowerPointReader(
|
|
chunk_size=500,
|
|
split_by="sentence",
|
|
include_image=False,
|
|
separate_table=False,
|
|
slide_prefix=None,
|
|
slide_suffix=None,
|
|
)
|
|
ppt_path = os.path.join(
|
|
os.path.abspath(os.path.dirname(__file__)),
|
|
"../tests/test.pptx",
|
|
)
|
|
docs = await reader(ppt_path=ppt_path)
|
|
|
|
# Without slide_prefix/suffix, content should not have XML tags
|
|
doc_texts = [_.metadata.content.get("text") for _ in docs]
|
|
|
|
# Verify exact content without slide tags
|
|
self.assertEqual(
|
|
doc_texts[0],
|
|
"AgentScope\nText content in slide 1",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[1],
|
|
"Title 2\n"
|
|
"Text content above table\n"
|
|
"| Name | Age | Career |\n"
|
|
"| --- | --- | --- |\n"
|
|
"| Alice | 25 | Teacher |\n"
|
|
"| Bob | 26 | Doctor |\n"
|
|
"\n"
|
|
"Text content below table",
|
|
)
|
|
self.assertEqual(
|
|
doc_texts[2],
|
|
"Title 3\ntext content above image\ntext content below image",
|
|
)
|