From b2f0b9f9f13cac11af1e95f75654b4fa3c112837 Mon Sep 17 00:00:00 2001 From: Ryan Chen Date: Mon, 12 Oct 2026 17:52:24 -0400 Subject: [PATCH] Add conversation file upload and file_key in stream-query MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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//files/. 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 and records 📎 name on the saved message. Co-Authored-By: Claude Opus 5.5 --- blueprints/conversation/__init__.py | 104 +++++++++++- tests/unit/test_file_upload.py | 252 ++++++++++++++++++++++++++++ utils/file_extract.py | 62 +++++++ utils/s3_client.py | 25 +++ 4 files changed, 442 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_file_upload.py create mode 100644 utils/file_extract.py diff --git a/blueprints/conversation/__init__.py b/blueprints/conversation/__init__.py index a9ea890..d3bf8f7 100644 --- a/blueprints/conversation/__init__.py +++ b/blueprints/conversation/__init__.py @@ -5,6 +5,7 @@ import os import re import time import uuid +from urllib.parse import quote, unquote from quart import Blueprint, jsonify, make_response, request from quart_jwt_extended import get_jwt_identity @@ -16,9 +17,19 @@ from blueprints.users.decorators import jwt_access_required from metrics import count_result, observe_response, time_response from utils.image_process import image_content_block from utils.strip_markdown import extract_agent_text +from utils.file_extract import ( + FILE_TYPES, + FileExtractionError, + attachment_block, + extract_text, + file_extension, +) +from utils.image_upload import ALLOWED_TYPES as IMAGE_TYPES from utils.image_upload import ImageValidationError, process_image from utils.s3_client import generate_presigned_url as s3_presigned_url +from utils.s3_client import get_file as s3_get_file from utils.s3_client import get_image as s3_get_image +from utils.s3_client import upload_file as s3_upload_file from utils.s3_client import upload_image as s3_upload_image from .logic import ( @@ -203,6 +214,79 @@ async def upload_image(): return jsonify({"image_key": key}) +def _file_key_belongs_to(file_key: str, conversation_id) -> bool: + """Only allow attachments stored under this conversation's prefix.""" + return ( + isinstance(file_key, str) + and ".." not in file_key + and file_key.startswith(f"conversations/{conversation_id}/files/") + ) + + +async def _load_attached_file(file_key: str) -> tuple[str, str]: + """Fetch an attached file from S3 and return (original name, extracted text).""" + data, metadata = await s3_get_file(file_key) + name = unquote(metadata.get("filename", "")) or file_key.rsplit("/", 1)[-1] + ext = file_extension(file_key, None) or "" + return name, extract_text(data, ext) + + +@conversation_blueprint.post("/upload") +@jwt_access_required +async def upload_attachment(): + current_user_uuid = get_jwt_identity() + files = await request.files + form = await request.form + file = files.get("file") + conversation_id = form.get("conversation_id") + + if not file or not conversation_id: + return jsonify({"error": "file and conversation_id are required"}), 400 + + conversation = await _get_owned_conversation(conversation_id, current_user_uuid) + if conversation is None: + return jsonify({"error": "conversation not found"}), 404 + + file_bytes = file.read() + content_type = (file.content_type or "").split(";")[0].strip().lower() + name = os.path.basename(file.filename or "") or "attachment" + + if content_type in IMAGE_TYPES: + try: + processed_bytes, output_content_type = process_image( + file_bytes, content_type + ) + except Exception as e: + return jsonify({"error": f"Could not process image: {e}"}), 400 + ext = output_content_type.split("/")[-1] + if ext == "jpeg": + ext = "jpg" + key = f"conversations/{conversation.id}/{uuid.uuid4()}.{ext}" + await s3_upload_image(processed_bytes, key, output_content_type) + return jsonify({"kind": "image", "image_key": key}) + + ext = file_extension(name, content_type) + if ext is None: + return ( + jsonify( + { + "error": f"Unsupported file type: {content_type or name}. " + "Allowed: images, PDF, TXT, MD, CSV" + } + ), + 415, + ) + + try: + text = extract_text(file_bytes, ext) + except FileExtractionError as e: + return jsonify({"error": str(e)}), 400 + + key = f"conversations/{conversation.id}/files/{uuid.uuid4()}.{ext}" + await s3_upload_file(file_bytes, key, FILE_TYPES[ext], {"filename": quote(name)}) + return jsonify({"kind": "file", "file_key": key, "name": name, "chars": len(text)}) + + @conversation_blueprint.get("/image/") @jwt_access_required async def serve_image(image_key: str): @@ -219,12 +303,30 @@ async def stream_query(): query_text = data.get("query") conversation_id = data.get("conversation_id") image_key = data.get("image_key") + file_key = data.get("file_key") conversation = await _get_owned_conversation(conversation_id, current_user_uuid) if conversation is None: return jsonify({"error": "conversation not found"}), 404 + + # An attached file is extracted and appended to the agent's user turn; the + # saved message only records its name. + saved_text = query_text or "" + if file_key: + if not _file_key_belongs_to(file_key, conversation.id): + return jsonify({"error": "invalid file_key"}), 400 + try: + file_name, file_text = await _load_attached_file(file_key) + block = attachment_block(file_name, file_text) + except Exception as e: + logging.error(f"Failed to load file {file_key}: {e}") + file_name = file_key.rsplit("/", 1)[-1] + block = "[A file was attached but could not be loaded.]" + query_text = f"{query_text or 'The user attached this file.'}\n\n{block}" + saved_text = f"{saved_text}\n\n📎 {file_name}".strip() + await add_message_to_conversation( conversation=conversation, - message=query_text or "", + message=saved_text, speaker="user", user=user, image_key=image_key, diff --git a/tests/unit/test_file_upload.py b/tests/unit/test_file_upload.py new file mode 100644 index 0000000..bf18dae --- /dev/null +++ b/tests/unit/test_file_upload.py @@ -0,0 +1,252 @@ +"""Tests for conversation file attachments: extraction, /upload and stream-query file_key.""" + +import io +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import fitz +import pytest +from PIL import Image +from quart import Quart +from quart.datastructures import FileStorage +from quart_jwt_extended import JWTManager, create_access_token + +import blueprints.conversation as conv +from utils.file_extract import ( + MAX_ATTACHMENT_CHARS, + TRUNCATION_MARKER, + FileExtractionError, + attachment_block, + cap_text, + extract_text, + file_extension, +) + +CONV_ID = "11111111-1111-1111-1111-111111111111" +OTHER_ID = "22222222-2222-2222-2222-222222222222" + + +def _pdf_bytes(text: str) -> bytes: + doc = fitz.open() + page = doc.new_page() + page.insert_text((72, 72), text) + data = doc.tobytes() + doc.close() + return data + + +def _png_bytes() -> bytes: + buf = io.BytesIO() + Image.new("RGB", (4, 4), "orange").save(buf, format="PNG") + return buf.getvalue() + + +class TestFileExtension: + @pytest.mark.parametrize( + "name,mime,expected", + [ + ("a.pdf", "application/pdf", "pdf"), + ("notes", "text/markdown", "md"), + ("notes.md", "application/octet-stream", "md"), + ("data.CSV", "", "csv"), + ("x.txt", "text/plain; charset=utf-8", "txt"), + ("x.markdown", None, "md"), + ("x.docx", "application/vnd.openxmlformats", None), + ("x.zip", "application/zip", None), + ], + ) + def test_detects_supported_types(self, name, mime, expected): + assert file_extension(name, mime) == expected + + +class TestExtractText: + def test_pdf(self): + assert "Simba eats kibble" in extract_text( + _pdf_bytes("Simba eats kibble"), "pdf" + ) + + def test_corrupt_pdf_raises(self): + with pytest.raises(FileExtractionError): + extract_text(b"not a pdf", "pdf") + + def test_text_decodes_utf8_with_replacement(self): + assert extract_text("héllo".encode() + b"\xff", "txt") == "héllo�" + + def test_csv_and_md_are_plain_text(self): + assert extract_text(b"\xef\xbb\xbfa,b\n1,2", "csv") == "a,b\n1,2" + assert extract_text(b"# Title", "md") == "# Title" + + +class TestTruncation: + def test_short_text_untouched(self): + assert cap_text("abc") == "abc" + + def test_long_text_capped_with_marker(self): + out = cap_text("x" * (MAX_ATTACHMENT_CHARS + 50)) + assert out.endswith(TRUNCATION_MARKER) + assert out.count("x") == MAX_ATTACHMENT_CHARS + + def test_attachment_block_wraps_and_caps(self): + block = attachment_block('a"b.txt', "y" * (MAX_ATTACHMENT_CHARS + 1)) + assert block.startswith('\n') + assert block.endswith(f"{TRUNCATION_MARKER}\n") + + +class TestFileKeyPrefix: + def test_own_conversation(self): + assert conv._file_key_belongs_to( + f"conversations/{CONV_ID}/files/x.pdf", CONV_ID + ) + + @pytest.mark.parametrize( + "key", + [ + f"conversations/{OTHER_ID}/files/x.pdf", + f"conversations/{CONV_ID}/x.jpg", + f"conversations/{CONV_ID}/files/../../{OTHER_ID}/files/x.pdf", + f"conversations/{CONV_ID}", + None, + ], + ) + def test_rejects_other_keys(self, key): + assert not conv._file_key_belongs_to(key, CONV_ID) + + +@pytest.fixture +def app(): + app = Quart(__name__) + app.config["JWT_SECRET_KEY"] = "test-secret" + app.config["JWT_ACCESS_TOKEN_EXPIRES"] = timedelta(hours=1) + JWTManager(app) + app.register_blueprint(conv.conversation_blueprint) + return app + + +@pytest.fixture +async def headers(app): + async with app.app_context(): + token = create_access_token(identity="user-1") + return {"Authorization": f"Bearer {token}"} + + +def _owned(conversation_id, user_id): + if conversation_id == CONV_ID and user_id == "user-1": + return SimpleNamespace(id=CONV_ID) + return None + + +@pytest.fixture +def s3(): + with ( + patch.object(conv, "_get_owned_conversation", AsyncMock(side_effect=_owned)), + patch.object(conv, "s3_upload_file", AsyncMock()) as upload_file, + patch.object(conv, "s3_upload_image", AsyncMock()) as upload_image, + ): + yield SimpleNamespace(upload_file=upload_file, upload_image=upload_image) + + +async def _upload( + app, headers, data: bytes, filename: str, mime: str, conversation_id=CONV_ID +): + return await app.test_client().post( + "/api/conversation/upload", + headers=headers, + form={"conversation_id": conversation_id}, + files={ + "file": FileStorage(io.BytesIO(data), filename=filename, content_type=mime) + }, + ) + + +class TestUploadRoute: + async def test_pdf_is_stored_with_filename(self, app, headers, s3): + resp = await _upload( + app, headers, _pdf_bytes("hello"), "Vet bill.pdf", "application/pdf" + ) + assert resp.status_code == 200 + body = await resp.get_json() + assert body["kind"] == "file" + assert body["name"] == "Vet bill.pdf" + assert body["chars"] >= len("hello") + assert body["file_key"].startswith(f"conversations/{CONV_ID}/files/") + assert body["file_key"].endswith(".pdf") + _, key, mime, metadata = s3.upload_file.await_args.args + assert key == body["file_key"] and mime == "application/pdf" + assert metadata == {"filename": "Vet%20bill.pdf"} + + async def test_markdown_by_extension(self, app, headers, s3): + resp = await _upload(app, headers, b"# hi", "n.md", "application/octet-stream") + body = await resp.get_json() + assert resp.status_code == 200 + assert body["chars"] == 4 and body["file_key"].endswith(".md") + + async def test_image_goes_through_process_image(self, app, headers, s3): + resp = await _upload(app, headers, _png_bytes(), "cat.png", "image/png") + assert resp.status_code == 200 + body = await resp.get_json() + assert body["kind"] == "image" + assert body["image_key"].startswith(f"conversations/{CONV_ID}/") + s3.upload_image.assert_awaited_once() + s3.upload_file.assert_not_awaited() + + async def test_unsupported_type_is_415(self, app, headers, s3): + resp = await _upload(app, headers, b"PK\x03\x04", "a.zip", "application/zip") + assert resp.status_code == 415 + s3.upload_file.assert_not_awaited() + + async def test_unsupported_image_type_is_415(self, app, headers, s3): + resp = await _upload(app, headers, b"GIF89a", "a.gif", "image/gif") + assert resp.status_code == 415 + + async def test_corrupt_pdf_is_400(self, app, headers, s3): + resp = await _upload(app, headers, b"nope", "a.pdf", "application/pdf") + assert resp.status_code == 400 + + async def test_other_users_conversation_is_404(self, app, headers, s3): + resp = await _upload( + app, headers, b"hi", "a.txt", "text/plain", conversation_id=OTHER_ID + ) + assert resp.status_code == 404 + s3.upload_file.assert_not_awaited() + + async def test_requires_auth(self, app, s3): + resp = await _upload(app, {}, b"hi", "a.txt", "text/plain") + assert resp.status_code == 401 + + +class TestStreamQueryFileKey: + async def test_rejects_file_key_from_other_conversation(self, app, headers, s3): + add = AsyncMock() + with ( + patch.object(conv.blueprints.users.models.User, "get", AsyncMock()), + patch.object(conv, "add_message_to_conversation", add), + ): + resp = await app.test_client().post( + "/api/conversation/stream-query", + headers=headers, + json={ + "query": "summarise", + "conversation_id": CONV_ID, + "file_key": f"conversations/{OTHER_ID}/files/x.pdf", + }, + ) + assert resp.status_code == 400 + add.assert_not_awaited() + + async def test_load_attached_file_recovers_name_and_text(self): + get = AsyncMock(return_value=(b"a,b\n1,2", {"filename": "M%C3%A9nage.csv"})) + with patch.object(conv, "s3_get_file", get): + name, text = await conv._load_attached_file( + f"conversations/{CONV_ID}/files/u.csv" + ) + assert name == "Ménage.csv" + assert text == "a,b\n1,2" + + async def test_load_attached_file_falls_back_to_key_name(self): + get = AsyncMock(return_value=(b"x", {})) + with patch.object(conv, "s3_get_file", get): + name, _ = await conv._load_attached_file( + f"conversations/{CONV_ID}/files/u.txt" + ) + assert name == "u.txt" diff --git a/utils/file_extract.py b/utils/file_extract.py new file mode 100644 index 0000000..9ab738b --- /dev/null +++ b/utils/file_extract.py @@ -0,0 +1,62 @@ +"""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'\n{cap_text(text)}\n' diff --git a/utils/s3_client.py b/utils/s3_client.py index 4d941a5..668ee75 100644 --- a/utils/s3_client.py +++ b/utils/s3_client.py @@ -61,3 +61,28 @@ async def delete_image(key: str) -> None: async with _get_client() as client: await client.delete_object(Bucket=S3_BUCKET_NAME, Key=key) logging.info(f"Deleted image from S3: {key}") + + +async def upload_file( + file_bytes: bytes, key: str, content_type: str, metadata: dict[str, str] +) -> str: + """Store a non-image object. Metadata values must be ASCII (S3 headers).""" + async with _get_client() as client: + await client.put_object( + Bucket=S3_BUCKET_NAME, + Key=key, + Body=file_bytes, + ContentType=content_type, + Metadata=metadata, + ) + logging.info(f"Uploaded file to S3: {key}") + return key + + +async def get_file(key: str) -> tuple[bytes, dict[str, str]]: + """Return an object's bytes and its user metadata.""" + async with _get_client() as client: + response = await client.get_object(Bucket=S3_BUCKET_NAME, Key=key) + body = await response["Body"].read() + metadata = response.get("Metadata", {}) or {} + return body, metadata