mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-29 15:23:51 +00:00
fix: handle encoding errors when adding files to vector store
- Add try-catch block around data.decode() to handle UnicodeDecodeError - Implement UTF-8 fallback when detected encoding fails - re-raise origional exception if fallback fails - add unit tests Fixes #2572: UnicodeDecodeError when uploading files with problematic encodings Signed-off-by: Derek Higgins <derekh@redhat.com>
This commit is contained in:
parent
b246b0660e
commit
87b8530e3f
2 changed files with 57 additions and 2 deletions
|
@ -92,7 +92,20 @@ def content_from_data_and_mime_type(data: bytes | str, mime_type: str | None, en
|
||||||
mime_category = mime_type.split("/")[0] if mime_type else None
|
mime_category = mime_type.split("/")[0] if mime_type else None
|
||||||
if mime_category == "text":
|
if mime_category == "text":
|
||||||
# For text-based files (including CSV, MD)
|
# For text-based files (including CSV, MD)
|
||||||
return data.decode(encoding)
|
encodings_to_try = [encoding]
|
||||||
|
if encoding != "utf-8":
|
||||||
|
encodings_to_try.append("utf-8")
|
||||||
|
first_exception = None
|
||||||
|
for encoding in encodings_to_try:
|
||||||
|
try:
|
||||||
|
return data.decode(encoding)
|
||||||
|
except UnicodeDecodeError as e:
|
||||||
|
if first_exception is None:
|
||||||
|
first_exception = e
|
||||||
|
log.warning(f"Decoding failed with {encoding}: {e}")
|
||||||
|
# raise the origional exception, if we got here there was at least 1 exception
|
||||||
|
log.error(f"Could not decode data as any of {encodings_to_try}")
|
||||||
|
raise first_exception
|
||||||
|
|
||||||
elif mime_type == "application/pdf":
|
elif mime_type == "application/pdf":
|
||||||
return parse_pdf(data)
|
return parse_pdf(data)
|
||||||
|
|
|
@ -10,7 +10,7 @@ import pytest
|
||||||
|
|
||||||
from llama_stack.apis.common.content_types import URL, TextContentItem
|
from llama_stack.apis.common.content_types import URL, TextContentItem
|
||||||
from llama_stack.apis.tools import RAGDocument
|
from llama_stack.apis.tools import RAGDocument
|
||||||
from llama_stack.providers.utils.memory.vector_store import content_from_doc
|
from llama_stack.providers.utils.memory.vector_store import content_from_data_and_mime_type, content_from_doc
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
|
@ -143,3 +143,45 @@ async def test_content_from_doc_with_interleaved_content():
|
||||||
|
|
||||||
assert result == "First item\nSecond item"
|
assert result == "First item\nSecond item"
|
||||||
mock_interleaved.assert_called_once_with(interleaved_content)
|
mock_interleaved.assert_called_once_with(interleaved_content)
|
||||||
|
|
||||||
|
|
||||||
|
def test_content_from_data_and_mime_type_success_utf8():
|
||||||
|
"""Test successful decoding with UTF-8 encoding."""
|
||||||
|
data = "Hello World! 🌍".encode()
|
||||||
|
mime_type = "text/plain"
|
||||||
|
|
||||||
|
with patch("chardet.detect") as mock_detect:
|
||||||
|
mock_detect.return_value = {"encoding": "utf-8"}
|
||||||
|
|
||||||
|
result = content_from_data_and_mime_type(data, mime_type)
|
||||||
|
|
||||||
|
mock_detect.assert_called_once_with(data)
|
||||||
|
assert result == "Hello World! 🌍"
|
||||||
|
|
||||||
|
|
||||||
|
def test_content_from_data_and_mime_type_error_win1252():
|
||||||
|
"""Test fallback to UTF-8 when Windows-1252 encoding detection fails."""
|
||||||
|
data = "Hello World! 🌍".encode()
|
||||||
|
mime_type = "text/plain"
|
||||||
|
|
||||||
|
with patch("chardet.detect") as mock_detect:
|
||||||
|
mock_detect.return_value = {"encoding": "Windows-1252"}
|
||||||
|
|
||||||
|
result = content_from_data_and_mime_type(data, mime_type)
|
||||||
|
|
||||||
|
assert result == "Hello World! 🌍"
|
||||||
|
mock_detect.assert_called_once_with(data)
|
||||||
|
|
||||||
|
|
||||||
|
def test_content_from_data_and_mime_type_both_encodings_fail():
|
||||||
|
"""Test that exceptions are raised when both primary and UTF-8 encodings fail."""
|
||||||
|
# Create invalid byte sequence that fails with both encodings
|
||||||
|
data = b"\xff\xfe\x00\x8f" # Invalid UTF-8 sequence
|
||||||
|
mime_type = "text/plain"
|
||||||
|
|
||||||
|
with patch("chardet.detect") as mock_detect:
|
||||||
|
mock_detect.return_value = {"encoding": "windows-1252"}
|
||||||
|
|
||||||
|
# Should raise an exception instead of returning empty string
|
||||||
|
with pytest.raises(UnicodeDecodeError):
|
||||||
|
content_from_data_and_mime_type(data, mime_type)
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue