Spaces:
Sleeping
Sleeping
File size: 3,455 Bytes
cfd3735 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
"""Module to test generic parsers."""
from typing import Iterator
import pytest
from langchain.document_loaders.base import BaseBlobParser
from langchain.document_loaders.blob_loaders import Blob
from langchain.document_loaders.parsers.generic import MimeTypeBasedParser
from langchain.schema import Document
class TestMimeBasedParser:
"""Test mime based parser."""
def test_without_fallback_parser(self) -> None:
class FirstCharParser(BaseBlobParser):
def lazy_parse(self, blob: Blob) -> Iterator[Document]:
"""Extract the first character of a blob."""
yield Document(page_content=blob.as_string()[0])
class SecondCharParser(BaseBlobParser):
def lazy_parse(self, blob: Blob) -> Iterator[Document]:
"""Extract the second character of a blob."""
yield Document(page_content=blob.as_string()[1])
parser = MimeTypeBasedParser(
handlers={
"text/plain": FirstCharParser(),
"text/html": SecondCharParser(),
},
)
blob = Blob(data=b"Hello World", mimetype="text/plain")
docs = parser.parse(blob)
assert len(docs) == 1
doc = docs[0]
assert doc.page_content == "H"
# Check text/html handler.
blob = Blob(data=b"Hello World", mimetype="text/html")
docs = parser.parse(blob)
assert len(docs) == 1
doc = docs[0]
assert doc.page_content == "e"
blob = Blob(data=b"Hello World", mimetype="text/csv")
with pytest.raises(ValueError, match="Unsupported mime type"):
# Check that the fallback parser is used when the mimetype is not found.
parser.parse(blob)
def test_with_fallback_parser(self) -> None:
class FirstCharParser(BaseBlobParser):
def lazy_parse(self, blob: Blob) -> Iterator[Document]:
"""Extract the first character of a blob."""
yield Document(page_content=blob.as_string()[0])
class SecondCharParser(BaseBlobParser):
def lazy_parse(self, blob: Blob) -> Iterator[Document]:
"""Extract the second character of a blob."""
yield Document(page_content=blob.as_string()[1])
class ThirdCharParser(BaseBlobParser):
def lazy_parse(self, blob: Blob) -> Iterator[Document]:
"""Extract the third character of a blob."""
yield Document(page_content=blob.as_string()[2])
parser = MimeTypeBasedParser(
handlers={
"text/plain": FirstCharParser(),
"text/html": SecondCharParser(),
},
fallback_parser=ThirdCharParser(),
)
blob = Blob(data=b"Hello World", mimetype="text/plain")
docs = parser.parse(blob)
assert len(docs) == 1
doc = docs[0]
assert doc.page_content == "H"
# Check text/html handler.
blob = Blob(data=b"Hello World", mimetype="text/html")
docs = parser.parse(blob)
assert len(docs) == 1
doc = docs[0]
assert doc.page_content == "e"
# Check that the fallback parser is used when the mimetype is not found.
blob = Blob(data=b"Hello World", mimetype="text/csv")
docs = parser.parse(blob)
assert len(docs) == 1
doc = docs[0]
assert doc.page_content == "l"
|