mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-08-05 02:45:28 +00:00
Merge 6865cacaf2 into 20e7fc0164
This commit is contained in:
commit
ca9b8a0705
4 changed files with 189 additions and 23 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
3
uv.lock
generated
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
version = 1
|
||||
revision = 3
|
||||
requires-python = ">=3.14"
|
||||
Loading…
Add table
Reference in a new issue