ryan
b2f0b9f9f1
POST /api/conversation/upload takes an image (via process_image) or a PDF/txt/md/csv file, checks the conversation belongs to the caller, and stores files at conversations/<id>/files/<uuid>.<ext> with the original name in S3 metadata. Anything else gets a 415. stream-query accepts file_key (must sit under this conversation's files/ prefix), extracts its text capped at 12k chars, appends it to the agent's user turn as <attached file> and records 📎 name on the saved message. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
"""Text extraction for files attached to a conversation (PDF, txt, md, csv)."""
|
|
|
|
import os
|
|
|
|
import fitz
|
|
|
|
# Extension -> canonical MIME type for the text-like files we accept.
|
|
FILE_TYPES = {
|
|
"pdf": "application/pdf",
|
|
"txt": "text/plain",
|
|
"md": "text/markdown",
|
|
"csv": "text/csv",
|
|
}
|
|
_MIME_TO_EXT = {mime: ext for ext, mime in FILE_TYPES.items()}
|
|
_MIME_TO_EXT["text/x-markdown"] = "md"
|
|
|
|
MAX_ATTACHMENT_CHARS = 12_000
|
|
TRUNCATION_MARKER = "[truncated]"
|
|
|
|
|
|
class FileExtractionError(Exception):
|
|
pass
|
|
|
|
|
|
def file_extension(filename: str | None, content_type: str | None) -> str | None:
|
|
"""Return the supported extension for an upload, or None if unsupported.
|
|
|
|
Matches on the MIME type first, then falls back to the filename extension
|
|
(browsers often send .md and .csv as application/octet-stream).
|
|
"""
|
|
mime = (content_type or "").split(";")[0].strip().lower()
|
|
if mime in _MIME_TO_EXT:
|
|
return _MIME_TO_EXT[mime]
|
|
ext = os.path.splitext(filename or "")[1].lstrip(".").lower()
|
|
if ext == "markdown":
|
|
ext = "md"
|
|
return ext if ext in FILE_TYPES else None
|
|
|
|
|
|
def extract_text(data: bytes, ext: str) -> str:
|
|
"""Extract plain text from file bytes. Text files decode as UTF-8 with replacement."""
|
|
if ext == "pdf":
|
|
try:
|
|
with fitz.open(stream=data, filetype="pdf") as doc:
|
|
return "\n".join(page.get_text() for page in doc).strip()
|
|
except Exception as e:
|
|
raise FileExtractionError(f"Could not read PDF: {e}") from e
|
|
if ext in FILE_TYPES:
|
|
return data.decode("utf-8", errors="replace").lstrip("")
|
|
raise FileExtractionError(f"Unsupported file type: {ext}")
|
|
|
|
|
|
def cap_text(text: str, limit: int = MAX_ATTACHMENT_CHARS) -> str:
|
|
if len(text) <= limit:
|
|
return text
|
|
return f"{text[:limit]}\n{TRUNCATION_MARKER}"
|
|
|
|
|
|
def attachment_block(name: str, text: str) -> str:
|
|
"""Wrap (capped) file text for the user turn sent to the agent."""
|
|
safe_name = name.replace('"', "'")
|
|
return f'<attached file "{safe_name}">\n{cap_text(text)}\n</attached>'
|