This commit is contained in:
Matt Van Horn 2026-08-04 06:46:30 -07:00 committed by GitHub
commit ca9b8a0705
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 189 additions and 23 deletions

View file

@ -9,6 +9,8 @@ import tempfile
from typing import List, Dict, Any
from src.llm_core import llm_call
from src.markitdown_runtime import is_markitdown_format
from src.upload_handler import is_text_attachment
logger = logging.getLogger(__name__)
@ -16,28 +18,25 @@ MAX_INLINE_ATTACHMENT_CHARS = 24000
MIN_INLINE_ATTACHMENT_SLICE = 500
def _is_text_file(path: str) -> bool:
"""Check if file has text extension."""
return any(
path.lower().endswith(ext)
for ext in (".txt", ".py", ".html", ".htm", ".md", ".json", ".csv", ".log", ".js", ".nix")
)
def _is_text_file(path: str, content_type: str | None = None) -> bool:
"""Check a file using the shared safe text-attachment classification."""
return is_text_attachment(path, content_type)
def _process_text_file(path: str) -> str:
def _process_text_file(path: str, display_name: str | None = None) -> str:
"""Process text file with enhanced formatting and metadata."""
language_map = {
".py": "python", ".js": "javascript", ".html": "html", ".css": "css",
".py": "python", ".js": "javascript", ".html": "html", ".htm": "html", ".css": "css",
".json": "json", ".md": "markdown", ".txt": "text", ".csv": "csv",
".log": "log", ".sh": "bash", ".bash": "bash", ".nix": "nix",
".yml": "yaml", ".yaml": "yaml",
".xml": "xml", ".sql": "sql", ".cpp": "cpp", ".c": "c",
".xml": "xml", ".sql": "sql", ".cpp": "cpp", ".c": "c", ".h": "c",
".java": "java", ".go": "go", ".rs": "rust", ".php": "php",
".rb": "ruby", ".ts": "typescript", ".jsx": "javascript", ".tsx": "typescript",
}
filename = os.path.basename(path)
_, ext = os.path.splitext(path.lower())
filename = os.path.basename(display_name or path)
_, ext = os.path.splitext(filename.lower())
language = language_map.get(ext, "text")
max_len = 30000 if ext != ".log" else 10000
@ -92,9 +91,9 @@ def _process_text_file(path: str) -> str:
header += f"[Type: {language}, Lines: {line_count}, Size: {size_str} bytes]"
code_extensions = {
".py", ".js", ".html", ".css", ".json", ".md", ".sh", ".bash", ".nix",
".py", ".js", ".html", ".htm", ".css", ".json", ".md", ".sh", ".bash", ".nix",
".yml", ".yaml", ".xml", ".sql", ".cpp", ".c", ".java", ".go", ".rs", ".php", ".rb",
".ts", ".jsx", ".tsx",
".ts", ".jsx", ".tsx", ".h",
}
if ext in code_extensions:
code_block = f"```{language}\n{content}"
@ -435,6 +434,7 @@ def build_user_content(
_, ext = os.path.splitext(path.lower())
mime = upload_info.get("mime") or mimetypes.guess_type(path)[0] or "application/octet-stream"
display_name = upload_info.get("name") or upload_info.get("original_name") or path
_, display_ext = os.path.splitext(display_name.lower())
if upload_handler.is_image_file(display_name, mime):
try:
@ -472,7 +472,7 @@ def build_user_content(
content.insert(0, {"type": "text", "text": "[Audio attached but could not be processed]"})
elif upload_handler.is_document_file(display_name, mime):
if mime == "application/pdf":
if mime.partition(";")[0].strip().lower() == "application/pdf" or display_ext == ".pdf":
extracted_text = None
if session_id:
try:
@ -570,8 +570,8 @@ def build_user_content(
logger.warning(f"PDF auto-doc creation failed for {path}: {e}")
if extracted_text is None:
extracted_text = _process_pdf(path, owner=owner)
elif mime.startswith("text/") or _is_text_file(path):
extracted_text = _process_text_file(path)
elif _is_text_file(display_name, mime) and not is_markitdown_format(display_name):
extracted_text = _process_text_file(path, display_name=display_name)
else:
extracted_text = _process_office_document(
path,

View file

@ -64,6 +64,20 @@ ATTACHMENT_REFERENCE_LINE_RE = re.compile(
re.IGNORECASE,
)
TEXT_ATTACHMENT_EXTENSIONS = frozenset({
".txt", ".py", ".js", ".html", ".htm", ".css", ".json", ".md",
".csv", ".log", ".xml", ".yml", ".yaml", ".nix", ".sql", ".sh",
".bash", ".c", ".cpp", ".h", ".java", ".go", ".rs", ".php", ".rb",
".ts", ".jsx", ".tsx",
})
def is_text_attachment(filename: str, content_type: str | None = None) -> bool:
"""Return True for supported text/code extensions or a ``text/*`` MIME."""
_, ext = os.path.splitext((filename or "").lower())
mime = (content_type or "").partition(";")[0].strip().lower()
return ext in TEXT_ATTACHMENT_EXTENSIONS or mime.startswith("text/")
def is_valid_upload_id(upload_id: str) -> bool:
"""Return True when *upload_id* matches the canonical uploads.json id format."""
@ -309,10 +323,6 @@ class UploadHandler:
"""Check if a file is a document based on extension or content type."""
document_extensions = {
'.pdf', '.docx', '.xlsx', '.pptx', '.xls', '.epub',
'.txt', '.py', '.js', '.html', '.htm',
'.css', '.json', '.md', '.csv', '.log', '.xml', '.yml',
'.yaml', '.nix', '.sql', '.sh', '.bash', '.c', '.cpp', '.h',
'.java', '.go', '.rs', '.php', '.rb', '.ts', '.jsx', '.tsx'
}
document_mime_types = {
'application/pdf',
@ -321,8 +331,10 @@ class UploadHandler:
'application/vnd.openxmlformats-officedocument.presentationml.presentation',
'application/vnd.ms-excel',
'application/epub+zip',
'text/plain'
}
if is_text_attachment(filename, content_type):
return True
# Check by extension
_, ext = os.path.splitext(filename.lower())
@ -330,7 +342,8 @@ class UploadHandler:
return True
# Check by content type if provided
if content_type and content_type in document_mime_types:
mime = (content_type or '').partition(';')[0].strip().lower()
if mime in document_mime_types:
return True
return False

View file

@ -1,5 +1,8 @@
import pytest
import src.document_processor as document_processor
from src.document_processor import _is_text_file, _process_text_file
from src.upload_handler import UploadHandler
from src.upload_handler import UploadHandler, is_text_attachment
def test_nix_files_are_treated_as_readable_documents(tmp_path):
@ -18,3 +21,150 @@ def test_nix_file_processing_includes_content_in_code_block(tmp_path):
assert "[Type: nix" in rendered
assert "```nix" in rendered
assert "services.openssh.enable = true;" in rendered
def _build_user_content(tmp_path, *, stored_name, display_name, mime, body):
upload_dir = tmp_path / "uploads"
upload_dir.mkdir(exist_ok=True)
path = upload_dir / stored_name
path.write_bytes(body)
handler = UploadHandler(str(tmp_path), str(upload_dir))
return document_processor.build_user_content(
"Read this attachment.",
["attachment-id"],
str(upload_dir),
handler,
owner="tester",
resolved_uploads={
"attachment-id": {
"path": str(path),
"name": display_name,
"mime": mime,
}
},
)
def test_text_mime_without_recognized_extension_reaches_user_content(tmp_path):
handler = UploadHandler(str(tmp_path), str(tmp_path / "uploads"))
mime = "Text/Markdown; charset=utf-8"
assert is_text_attachment("release-notes.unknown", mime)
assert handler.is_document_file("release-notes.unknown", mime)
content = _build_user_content(
tmp_path,
stored_name="0123456789abcdef0123456789abcdef",
display_name="release-notes.unknown",
mime=mime,
body=b"Deployment requires a database backup.\n",
)
assert "=== File: release-notes.unknown ===" in content
assert "Deployment requires a database backup." in content
@pytest.mark.parametrize(
("filename", "body"),
[
("include/widget.h", b"#define WIDGET_LIMIT 8\n"),
("deploy/config.yaml", b"replicas: 3\n"),
],
)
def test_text_extension_with_opaque_mime_reaches_user_content(tmp_path, filename, body):
display_name = filename.rsplit("/", 1)[-1]
assert is_text_attachment(filename, "application/octet-stream")
content = _build_user_content(
tmp_path,
stored_name=f"0123456789abcdef0123456789abcdef.{display_name.rsplit('.', 1)[-1]}",
display_name=display_name,
mime="application/octet-stream",
body=body,
)
assert f"=== File: {display_name} ===" in content
assert body.decode().strip() in content
def test_pdf_and_office_files_are_not_classified_as_generic_text(tmp_path):
handler = UploadHandler(str(tmp_path), str(tmp_path / "uploads"))
assert handler.is_document_file("packet.pdf", "application/pdf")
assert handler.is_document_file(
"report.docx",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
)
assert not is_text_attachment("packet.pdf", "application/pdf")
assert not is_text_attachment(
"report.docx",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
)
def test_pdf_and_office_files_keep_specialized_processing_paths(tmp_path, monkeypatch):
calls = []
monkeypatch.setattr(
document_processor,
"_process_pdf",
lambda path, owner=None: calls.append(("pdf", path)) or "\n\n[pdf text]",
)
monkeypatch.setattr(
document_processor,
"_process_office_document",
lambda path, display_name, **kwargs: calls.append(("office", path)) or "\n\n[office text]",
)
pdf_content = _build_user_content(
tmp_path,
stored_name="0123456789abcdef0123456789abcdef.pdf",
display_name="packet.pdf",
mime="application/octet-stream",
body=b"%PDF-1.4 fake",
)
office_content = _build_user_content(
tmp_path,
stored_name="fedcba9876543210fedcba9876543210.docx",
display_name="report.docx",
mime="application/octet-stream",
body=b"fake office file",
)
assert [kind for kind, _path in calls] == ["pdf", "office"]
assert "[pdf text]" in pdf_content
assert "[office text]" in office_content
def test_office_extension_takes_precedence_over_text_mime(tmp_path, monkeypatch):
calls = []
monkeypatch.setattr(
document_processor,
"_process_office_document",
lambda path, display_name, **kwargs: calls.append(display_name) or "\n\n[office text]",
)
content = _build_user_content(
tmp_path,
stored_name="0123456789abcdef0123456789abcdef.docx",
display_name="report.docx",
mime="text/plain; charset=utf-8",
body=b"fake office file",
)
assert calls == ["report.docx"]
assert "[office text]" in content
def test_unknown_binary_attachment_is_not_decoded_as_text(tmp_path):
content = _build_user_content(
tmp_path,
stored_name="0123456789abcdef0123456789abcdef.bin",
display_name="payload.bin",
mime="application/octet-stream",
body=b"SECRET_ASCII_INSIDE_BINARY\x00\xff",
)
assert "[Attached non-text file]" in content
assert "SECRET_ASCII_INSIDE_BINARY" not in content

3
uv.lock generated Normal file
View file

@ -0,0 +1,3 @@
version = 1
revision = 3
requires-python = ">=3.14"