diff --git a/python/samples/05-end-to-end/chatkit-integration/app.py b/python/samples/05-end-to-end/chatkit-integration/app.py index 612d9f9c443..97097cc3945 100644 --- a/python/samples/05-end-to-end/chatkit-integration/app.py +++ b/python/samples/05-end-to-end/chatkit-integration/app.py @@ -43,7 +43,7 @@ # ChatKit imports from chatkit.actions import Action from chatkit.server import ChatKitServer -from chatkit.store import StoreItemType, default_generate_id +from chatkit.store import NotFoundError, StoreItemType, default_generate_id from chatkit.types import ( ThreadItem, ThreadItemDoneEvent, @@ -595,20 +595,25 @@ async def upload_file(attachment_id: str, file: UploadFile = File(...)): # noqa logger.warning(f"Rejected invalid attachment ID: {attachment_id!r}") return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."}) + try: + attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) + except NotFoundError: + return JSONResponse(status_code=404, content={"error": "Attachment not found."}) + + if attachment.upload_descriptor is None: + return JSONResponse(status_code=409, content={"error": "Attachment upload is already complete."}) + try: # Read file contents contents = await file.read() # Save to disk - file_path.write_bytes(contents) + file_path.write_bytes(contents) # CodeQL [SM01305] Path is constrained by get_file_path. logger.info(f"Saved {len(contents)} bytes to {file_path}") - # Load the attachment metadata from the data store - attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) - - # Clear the upload_url since upload is complete - attachment.upload_url = None + # Clear the upload descriptor since upload is complete + attachment.upload_descriptor = None # Save the updated attachment back to the store await data_store.save_attachment(attachment, {"user_id": DEFAULT_USER_ID}) @@ -637,19 +642,15 @@ async def preview_image(attachment_id: str): return JSONResponse(status_code=400, content={"error": "Invalid attachment ID."}) try: - if not file_path.exists(): - return JSONResponse(status_code=404, content={"error": "File not found"}) + attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) + except NotFoundError: + return JSONResponse(status_code=404, content={"error": "Attachment not found."}) - # Determine media type from file extension or attachment metadata - # For simplicity, we'll try to load from the store - try: - attachment = await data_store.load_attachment(attachment_id, {"user_id": DEFAULT_USER_ID}) - media_type = attachment.mime_type - except Exception: - # Default to binary if we can't determine - media_type = "application/octet-stream" + try: + if not file_path.exists(): # CodeQL [SM01305] Path is constrained by get_file_path. + return JSONResponse(status_code=404, content={"error": "File not found"}) - return FileResponse(file_path, media_type=media_type) + return FileResponse(file_path, media_type=attachment.mime_type) # CodeQL [SM01305] Path is constrained by get_file_path. # fmt: skip except Exception as e: logger.error(f"Error serving preview for attachment {attachment_id}: {e}", exc_info=True) diff --git a/python/samples/05-end-to-end/chatkit-integration/attachment_store.py b/python/samples/05-end-to-end/chatkit-integration/attachment_store.py index b08ae9c43a4..c4de028b29a 100644 --- a/python/samples/05-end-to-end/chatkit-integration/attachment_store.py +++ b/python/samples/05-end-to-end/chatkit-integration/attachment_store.py @@ -11,7 +11,13 @@ from typing import TYPE_CHECKING, Any from chatkit.store import AttachmentStore -from chatkit.types import Attachment, AttachmentCreateParams, FileAttachment, ImageAttachment +from chatkit.types import ( + Attachment, + AttachmentCreateParams, + AttachmentUploadDescriptor, + FileAttachment, + ImageAttachment, +) from pydantic import AnyUrl if TYPE_CHECKING: @@ -70,7 +76,7 @@ def get_file_path(self, attachment_id: str) -> Path: if not attachment_id or attachment_id in {".", ".."} or "/" in attachment_id or "\\" in attachment_id: raise ValueError(f"Invalid attachment ID: {attachment_id!r}") - file_path = (self.uploads_dir / attachment_id).resolve() + file_path = (self.uploads_dir / attachment_id).resolve() # CodeQL [SM01305] Path containment is validated below. # fmt: skip if not file_path.is_relative_to(self.uploads_dir) or file_path.parent != self.uploads_dir: raise ValueError(f"Invalid attachment ID: {attachment_id!r}") return file_path @@ -90,8 +96,11 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s # Generate unique ID for this attachment attachment_id = self.generate_attachment_id(input.mime_type, context) - # Generate upload URL that points to our FastAPI upload endpoint - upload_url = f"{self.base_url}/upload/{attachment_id}" + # Generate upload instructions that point to our FastAPI upload endpoint + upload_descriptor = AttachmentUploadDescriptor( + url=AnyUrl(f"{self.base_url}/upload/{attachment_id}"), + method="POST", + ) # Create appropriate attachment type based on MIME type if input.mime_type.startswith("image/"): @@ -103,7 +112,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s type="image", mime_type=input.mime_type, name=input.name, - upload_url=AnyUrl(upload_url), + upload_descriptor=upload_descriptor, preview_url=AnyUrl(preview_url), ) else: @@ -113,7 +122,7 @@ async def create_attachment(self, input: AttachmentCreateParams, context: dict[s type="file", mime_type=input.mime_type, name=input.name, - upload_url=AnyUrl(upload_url), + upload_descriptor=upload_descriptor, ) # Save attachment metadata to data store so it's available during upload