🧱(files) allow to use S3 storage without external access
Some architectures do not expose their S3, in such cases it is only available through the backend. This commit proposes two implementations to manage this: - frontend can now upload files to the backend (no direct access to S3) - two new modes to send file to the LLM: a temporary URL on the backend, or directly the file in b64.
This commit is contained in:
@@ -8,6 +8,10 @@ and this project adheres to
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- 🧱(files) allow to use S3 storage without external access #849
|
||||
|
||||
### Changed
|
||||
|
||||
- 🏗️(back) migrate to uv
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
# File Upload Modes
|
||||
|
||||
This document describes the different modes for handling file uploads in the Conversations application, and how to configure and use them.
|
||||
|
||||
## Overview
|
||||
|
||||
The application supports two independent configuration points:
|
||||
|
||||
1. **`FILE_UPLOAD_MODE`**: how the frontend uploads files (frontend → storage/backend)
|
||||
2. **`FILE_TO_LLM_MODE`**: how the backend provides files to the LLM (backend → LLM)
|
||||
|
||||
Each mode has different trade-offs in terms of security, performance, and LLM accessibility. The two settings can be combined based on your network constraints.
|
||||
|
||||
## Configuration
|
||||
|
||||
### Frontend upload mode (`FILE_UPLOAD_MODE`)
|
||||
|
||||
```bash
|
||||
# Default: presigned URL upload (backward compatible)
|
||||
FILE_UPLOAD_MODE=presigned_url
|
||||
|
||||
# Frontend uploads directly to backend
|
||||
FILE_UPLOAD_MODE=backend_to_s3
|
||||
```
|
||||
|
||||
### Backend delivery mode (`FILE_TO_LLM_MODE`)
|
||||
|
||||
```bash
|
||||
# Default: presigned URL mode (backward compatible)
|
||||
FILE_TO_LLM_MODE=presigned_url
|
||||
|
||||
# Backend provides base64-encoded data URLs
|
||||
FILE_TO_LLM_MODE=backend_base64
|
||||
|
||||
# Backend provides temporary URLs through the backend
|
||||
FILE_TO_LLM_MODE=backend_temporary_url
|
||||
```
|
||||
|
||||
Additional settings for backend temporary URL mode:
|
||||
|
||||
```bash
|
||||
# Base URL to reach backend
|
||||
FILE_BACKEND_URL="http://localhost:8071"
|
||||
|
||||
# Expiration time for temporary URLs (in seconds, default: 180 = 3 minutes)
|
||||
FILE_BACKEND_TEMPORARY_URL_EXPIRATION=180
|
||||
```
|
||||
|
||||
## Mode Details
|
||||
|
||||
### 1. Presigned URL Mode (Default)
|
||||
|
||||
**Frontend upload configuration:** `FILE_UPLOAD_MODE=presigned_url`
|
||||
|
||||
**Backend delivery configuration:** `FILE_TO_LLM_MODE=presigned_url`
|
||||
|
||||
**How it works:**
|
||||
- Frontend requests a presigned URL from the backend
|
||||
- Frontend uploads the file directly to S3 using the presigned URL
|
||||
- Frontend notifies the backend when upload is complete
|
||||
- Backend initiates malware detection
|
||||
- Backend returns presigned S3 URLs to the LLM
|
||||
|
||||
**Advantages:**
|
||||
- Files don't pass through the backend server (lower bandwidth usage)
|
||||
- Faster uploads for large files (direct to S3)
|
||||
- S3 handles the upload, no backend load
|
||||
- Backward compatible with existing frontend implementations
|
||||
|
||||
**Disadvantages:**
|
||||
- S3 bucket must be accessible from the frontend
|
||||
- Presigned URLs can be leaked if not handled carefully
|
||||
- Frontend needs to handle S3 credentials/configuration
|
||||
|
||||
**LLM Access:**
|
||||
- Images: Presigned S3 URLs with expiration (default: 3 minutes)
|
||||
- Documents: Presigned S3 URLs with expiration (default: 3 minutes)
|
||||
|
||||
**When to use:**
|
||||
- When frontend has direct access to S3
|
||||
- When you want to minimize backend load
|
||||
- When S3 is publicly accessible or accessible via VPN
|
||||
|
||||
|
||||
### 2. Backend Base64 Mode
|
||||
|
||||
**Frontend upload configuration:** `FILE_UPLOAD_MODE=backend_to_s3`
|
||||
|
||||
**Backend delivery configuration:** `FILE_TO_LLM_MODE=backend_base64`
|
||||
|
||||
**How it works:**
|
||||
- Frontend uploads the file directly to the backend
|
||||
- Backend stores the file on S3
|
||||
- Backend reads the file, encodes it as base64, and creates a data URL
|
||||
- LLM receives the file as a base64-encoded data URL
|
||||
|
||||
**Advantages:**
|
||||
- S3 can be private/internal (not accessible from frontend)
|
||||
- Files always go through the backend for validation
|
||||
- No presigned URLs to manage
|
||||
- Better control over file access
|
||||
- Data URLs work with all LLMs that support file content
|
||||
|
||||
**Disadvantages:**
|
||||
- Backend memory usage increases (entire file loaded for base64 encoding)
|
||||
- Slower for very large files (encoding overhead)
|
||||
- Increased bandwidth on backend
|
||||
- Data URLs can be very large in responses
|
||||
|
||||
**LLM Access:**
|
||||
- Images: Base64-encoded data URLs (format: `data:image/png;base64,...`)
|
||||
- Documents: Base64-encoded data URLs (format: `data:application/pdf;base64,...`)
|
||||
|
||||
**When to use:**
|
||||
- When S3 is not accessible from the frontend
|
||||
- When you want all file uploads to go through the backend
|
||||
- When the LLM supports base64-encoded data URLs
|
||||
- For smaller files (< 50MB)
|
||||
|
||||
|
||||
### 3. Backend Temporary URL Mode
|
||||
|
||||
**Frontend upload configuration:** `FILE_UPLOAD_MODE=backend_to_s3`
|
||||
|
||||
**Backend delivery configuration:** `FILE_TO_LLM_MODE=backend_temporary_url`
|
||||
|
||||
**How it works:**
|
||||
- Frontend uploads the file directly to the backend
|
||||
- Backend stores the file on S3
|
||||
- Backend generates a secure temporary access token stored in cache (TTL: 3 minutes by default)
|
||||
- Backend returns a temporary URL pointing to the backend's file-stream endpoint
|
||||
- LLM receives the temporary URL and accesses the file through the backend
|
||||
- Backend validates the token and streams the file content from S3 to the LLM
|
||||
|
||||
**Advantages:**
|
||||
- S3 can be private/internal (not accessible from frontend or LLM directly)
|
||||
- Files always go through the backend for validation and access control
|
||||
- LLM doesn't need direct access to S3
|
||||
- Tokens expire quickly (better security than long-lived presigned URLs)
|
||||
- No large data URL strings in memory or responses
|
||||
- Lower backend memory usage than base64 mode
|
||||
- Centralized file access control through the backend
|
||||
- Good balance between security and performance
|
||||
|
||||
**Disadvantages:**
|
||||
- LLM must be able to access the backend server
|
||||
- File streaming goes through the backend (adds some latency)
|
||||
- Time-limited access (token expires)
|
||||
|
||||
**LLM Access:**
|
||||
- Images: Temporary backend URLs with format `/api/v1.0/file-stream/{temporary_key}/` (token expiration: configurable, default: 3 minutes)
|
||||
- Documents: Temporary backend URLs with format `/api/v1.0/file-stream/{temporary_key}/` (token expiration: configurable, default: 3 minutes)
|
||||
|
||||
**When to use:**
|
||||
- When S3 is not accessible from the frontend or LLM
|
||||
- When you want backend control over uploads and file access
|
||||
- When you want time-limited access to files with centralized control
|
||||
- When you want the LLM to access files through the backend gateway
|
||||
- For large files (backend streams directly from S3 without loading entirely into memory)
|
||||
@@ -6,14 +6,103 @@ for the LLM to access them, and then reverting them back to local URLs when
|
||||
storing the messages in the database.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import mimetypes
|
||||
import secrets
|
||||
from typing import Dict, Iterable
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
from pydantic_ai import DocumentUrl, ImageUrl, ModelMessage, ModelRequest, UserPromptPart
|
||||
|
||||
from core.file_upload.enums import FileToLLMMode
|
||||
from core.file_upload.utils import generate_retrieve_policy
|
||||
|
||||
from chat.models import ChatConversation
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_temporary_url(key: str) -> str:
|
||||
"""
|
||||
Generate a temporary URL for accessing a file through the backend.
|
||||
|
||||
Instead of using S3 presigned URLs, this creates a temporary access key
|
||||
that's stored in cache (3 minutes TTL). The LLM accesses the file through
|
||||
a backend endpoint that validates the key and streams the file content.
|
||||
|
||||
This approach:
|
||||
- Works even when S3 is not accessible from the LLM
|
||||
- Provides better security (key is time-limited and single-use)
|
||||
- Allows the backend to control file access centrally
|
||||
|
||||
Args:
|
||||
key (str): The S3 object key where the file is stored.
|
||||
|
||||
Returns:
|
||||
str: A temporary URL with format: /api/v1.0/file-stream/{temporary_key}/
|
||||
"""
|
||||
# Generate a secure random key
|
||||
temporary_key = secrets.token_urlsafe(32)
|
||||
|
||||
# Store the S3 key in cache
|
||||
cache_key = f"file_access:{temporary_key}"
|
||||
cache.set(cache_key, key, timeout=settings.FILE_BACKEND_TEMPORARY_URL_EXPIRATION)
|
||||
|
||||
logger.info("Generated temporary file access key for S3 key: %s", key)
|
||||
|
||||
# Return the URL that the LLM will use to access the file
|
||||
return f"{settings.FILE_BACKEND_URL}/api/v1.0/file-stream/{temporary_key}/"
|
||||
|
||||
|
||||
def _get_file_url_for_llm(key: str, mode: str | None = None) -> str:
|
||||
"""
|
||||
Get the appropriate URL for the LLM to access a file based on the upload mode.
|
||||
|
||||
Args:
|
||||
key (str): The S3 object key where the file is stored.
|
||||
mode (str, optional): The upload mode. Defaults to FILE_TO_LLM_MODE setting.
|
||||
|
||||
Returns:
|
||||
str: The URL or data URL for the LLM to use.
|
||||
|
||||
Supported modes:
|
||||
- presigned_url: Returns a presigned S3 URL (default)
|
||||
- backend_temporary_url: Returns a presigned URL with shorter expiration
|
||||
- backend_base64: Returns a data URL with base64-encoded file content
|
||||
"""
|
||||
if mode is None:
|
||||
mode = settings.FILE_TO_LLM_MODE
|
||||
|
||||
if mode == FileToLLMMode.BACKEND_BASE64:
|
||||
# Read file from S3 and encode as base64 data URL
|
||||
try:
|
||||
with default_storage.open(key, "rb") as file:
|
||||
file_content = file.read()
|
||||
# Detect MIME type from file extension or default to octet-stream
|
||||
mime_type, _ = mimetypes.guess_type(key)
|
||||
if not mime_type:
|
||||
mime_type = "application/octet-stream"
|
||||
|
||||
# Create data URL
|
||||
b64_content = base64.b64encode(file_content).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{b64_content}"
|
||||
except Exception: # pylint: disable=broad-except
|
||||
# Fall back to presigned URL on error
|
||||
logger.exception(
|
||||
"Failed to read file for base64 encoding, falling back to presigned URL"
|
||||
)
|
||||
return generate_retrieve_policy(key)
|
||||
|
||||
elif mode == FileToLLMMode.BACKEND_TEMPORARY_URL:
|
||||
return generate_temporary_url(key)
|
||||
|
||||
# FileToLLMMode.PRESIGNED_URL or default
|
||||
return generate_retrieve_policy(key)
|
||||
|
||||
|
||||
def update_local_urls(
|
||||
conversation: ChatConversation,
|
||||
@@ -21,7 +110,9 @@ def update_local_urls(
|
||||
updated_url: Dict[str, str] | None = None,
|
||||
) -> Iterable[ImageUrl | DocumentUrl]:
|
||||
"""
|
||||
Replace local image or document URLs in the content list to use presigned S3 URLs.
|
||||
Replace local image or document URLs in the content list to use appropriate S3 URLs
|
||||
based on the configured FILE_TO_LLM_MODE.
|
||||
|
||||
⚠️Be careful, `media_contents` are replaced in place.
|
||||
|
||||
Args:
|
||||
@@ -31,7 +122,7 @@ def update_local_urls(
|
||||
mapping of original URLs to updated URLs.
|
||||
Returns:
|
||||
Iterable[ImageUrl | DocumentUrl]: Updated iterable of UserContent objects
|
||||
with presigned URLs.
|
||||
with appropriate S3 URLs based on the configured mode.
|
||||
"""
|
||||
# When images are stored locally, there is no host in the URL, so we can
|
||||
# just check if the URL starts, frontend adds a prefix `/media-key/` to the key.
|
||||
@@ -41,7 +132,9 @@ def update_local_urls(
|
||||
# Filter only ImageUrl contents
|
||||
media_contents = (c for c in contents if isinstance(c, (ImageUrl, DocumentUrl)))
|
||||
|
||||
# Replace URLs with presigned URLs
|
||||
# Replace URLs with appropriate S3 URLs based on mode
|
||||
upload_mode = settings.FILE_TO_LLM_MODE
|
||||
|
||||
for content in media_contents:
|
||||
idx = content.url.find(local_media_url_prefix)
|
||||
|
||||
@@ -57,7 +150,7 @@ def update_local_urls(
|
||||
# except if the user tampers with the conversation.
|
||||
continue
|
||||
|
||||
content.url = generate_retrieve_policy(key)
|
||||
content.url = _get_file_url_for_llm(key, upload_mode)
|
||||
if updated_url is not None:
|
||||
updated_url[content.url] = _initial_url
|
||||
|
||||
@@ -68,7 +161,7 @@ def update_history_local_urls(
|
||||
conversation: ChatConversation, messages: list[ModelMessage]
|
||||
) -> list[ModelMessage]:
|
||||
"""
|
||||
Replace local image/documents URLs in the message list to use presigned S3 URLs.
|
||||
Replace local image/documents URLs in the message list to use appropriate S3 URLs.
|
||||
|
||||
⚠️Be careful, `messages` are replaced in place.
|
||||
|
||||
@@ -79,7 +172,7 @@ def update_history_local_urls(
|
||||
Args:
|
||||
messages (list[ModelMessage]): List of ModelMessage objects.
|
||||
Returns:
|
||||
list[ModelMessage]: Updated list of ModelMessage objects with presigned URLs.
|
||||
list[ModelMessage]: Updated list of ModelMessage objects with appropriate S3 URLs.
|
||||
"""
|
||||
# Filter only ModelRequest messages
|
||||
requests = (msg for msg in messages if isinstance(msg, ModelRequest))
|
||||
|
||||
@@ -10,7 +10,7 @@ from django_pydantic_field.rest_framework import SchemaField # pylint: disable=
|
||||
from drf_spectacular.utils import extend_schema_field
|
||||
from rest_framework import serializers
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.enums import AttachmentStatus, FileUploadMode
|
||||
from core.file_upload.utils import generate_upload_policy
|
||||
|
||||
from chat import models
|
||||
@@ -180,7 +180,11 @@ class ChatConversationAttachmentSerializer(serializers.ModelSerializer):
|
||||
|
||||
|
||||
class CreateChatConversationAttachmentSerializer(serializers.ModelSerializer):
|
||||
"""Serializer for creating chat conversation attachments."""
|
||||
"""Serializer for creating chat conversation attachments.
|
||||
|
||||
For presigned_url mode: returns 'policy' field with presigned URL for direct S3 upload
|
||||
For backend modes: does not return 'policy' field (upload handled via backend endpoint)
|
||||
"""
|
||||
|
||||
policy = serializers.SerializerMethodField()
|
||||
uploaded_by = serializers.HiddenField(default=serializers.CurrentUserDefault())
|
||||
@@ -190,9 +194,15 @@ class CreateChatConversationAttachmentSerializer(serializers.ModelSerializer):
|
||||
model = models.ChatConversationAttachment
|
||||
fields = ["id", "key", "content_type", "file_name", "size", "policy", "uploaded_by"]
|
||||
|
||||
def get_policy(self, attachment) -> str:
|
||||
"""Return the policy to use if the item is a file."""
|
||||
return generate_upload_policy(attachment.key)
|
||||
def get_policy(self, attachment) -> str | None:
|
||||
"""Return the policy (presigned URL) only for presigned_url mode."""
|
||||
upload_mode = settings.FILE_UPLOAD_MODE
|
||||
|
||||
# Only return presigned URL in presigned_url mode
|
||||
if upload_mode == FileUploadMode.PRESIGNED_URL:
|
||||
return generate_upload_policy(attachment.key)
|
||||
|
||||
return None
|
||||
|
||||
def validate_size(self, size: Optional[int]) -> Optional[int]:
|
||||
"""Validate that the size is not greater than the maximum allowed size."""
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Tests for local_media_url_processors."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@@ -12,7 +13,10 @@ from pydantic_ai import (
|
||||
UserPromptPart,
|
||||
)
|
||||
|
||||
from core.file_upload.enums import FileToLLMMode
|
||||
|
||||
from chat.agents.local_media_url_processors import (
|
||||
_get_file_url_for_llm,
|
||||
update_history_local_urls,
|
||||
update_local_urls,
|
||||
)
|
||||
@@ -121,3 +125,549 @@ def test_update_history_local_urls_no_user_prompt_parts():
|
||||
result = update_history_local_urls(conversation, messages)
|
||||
assert result == messages
|
||||
mock.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_local_urls_uses_get_file_url_for_llm(mock_get_file_url):
|
||||
"""Test that update_local_urls uses _get_file_url_for_llm for mode-aware URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "mode-aware-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].url == "mode-aware-url"
|
||||
mock_get_file_url.assert_called_once()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_get_file_url_presigned_mode(mock_policy):
|
||||
"""Test URL generation in presigned_url mode."""
|
||||
|
||||
mock_policy.return_value = "presigned_s3_url"
|
||||
|
||||
url = _get_file_url_for_llm("test/key.pdf", FileToLLMMode.PRESIGNED_URL)
|
||||
|
||||
assert url == "presigned_s3_url"
|
||||
mock_policy.assert_called_once_with("test/key.pdf")
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_temporary_url")
|
||||
def test_get_file_url_backend_temporary_url_mode(mock_temp_url):
|
||||
"""Test URL generation in backend_temporary_url mode."""
|
||||
|
||||
mock_temp_url.return_value = "/api/v1.0/file-stream/token123/"
|
||||
|
||||
url = _get_file_url_for_llm("test/key.pdf", FileToLLMMode.BACKEND_TEMPORARY_URL)
|
||||
|
||||
assert url == "/api/v1.0/file-stream/token123/"
|
||||
mock_temp_url.assert_called_once_with("test/key.pdf")
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_get_file_url_backend_base64_mode(mock_storage):
|
||||
"""Test URL generation in backend_base64 mode."""
|
||||
|
||||
# Mock file content
|
||||
file_content = b"PDF binary content"
|
||||
mock_file = BytesIO(file_content)
|
||||
mock_storage.return_value.__enter__.return_value = mock_file
|
||||
|
||||
url = _get_file_url_for_llm("test/key.pdf", FileToLLMMode.BACKEND_BASE64)
|
||||
|
||||
# URL should be a data URL
|
||||
assert url.startswith("data:")
|
||||
assert "base64" in url
|
||||
mock_storage.assert_called_once()
|
||||
|
||||
|
||||
@patch(
|
||||
"chat.agents.local_media_url_processors.default_storage.open", side_effect=Exception("S3 error")
|
||||
)
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_get_file_url_backend_base64_fallback(mock_policy, _mock_storage):
|
||||
"""Test fallback to presigned URL when base64 encoding fails."""
|
||||
|
||||
mock_policy.return_value = "fallback_presigned_url"
|
||||
|
||||
url = _get_file_url_for_llm("test/key.pdf", FileToLLMMode.BACKEND_BASE64)
|
||||
|
||||
# Should fall back to presigned URL
|
||||
assert url == "fallback_presigned_url"
|
||||
mock_policy.assert_called_once()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_local_urls_multiple_images_with_modes(mock_get_file_url):
|
||||
"""Test handling multiple images with mode-aware URL generation."""
|
||||
conversation = ChatConversationFactory()
|
||||
|
||||
# Mock different URLs for different calls
|
||||
urls = ["url1", "url2", "url3"]
|
||||
mock_get_file_url.side_effect = urls
|
||||
|
||||
key1 = f"{conversation.pk}/image1.jpg"
|
||||
key2 = f"{conversation.pk}/image2.png"
|
||||
key3 = f"{conversation.pk}/document.pdf"
|
||||
|
||||
contents = [
|
||||
ImageUrl(url=f"/media-key/{key1}"),
|
||||
ImageUrl(url=f"/media-key/{key2}"),
|
||||
DocumentUrl(url=f"/media-key/{key3}"),
|
||||
]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == "url1"
|
||||
assert result[1].url == "url2"
|
||||
assert result[2].url == "url3"
|
||||
assert mock_get_file_url.call_count == 3
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_local_urls_mixed_external_and_local_urls(mock_get_file_url):
|
||||
"""Test handling of mixed external and local URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "mode-aware-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [
|
||||
ImageUrl(url=f"/media-key/{key}"), # Local URL - will be processed
|
||||
ImageUrl(url="https://external.com/image.jpg"), # External URL - kept as is
|
||||
ImageUrl(url="http://another.com/image.png"), # External URL - kept as is
|
||||
]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == "mode-aware-url"
|
||||
assert result[1].url == "https://external.com/image.jpg"
|
||||
assert result[2].url == "http://another.com/image.png"
|
||||
# Only one local URL was processed
|
||||
assert mock_get_file_url.call_count == 1
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_history_local_urls_with_mode_detection(mock_get_file_url):
|
||||
"""Test that update_history_local_urls uses mode detection for URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "mode-aware-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
user_prompt_content = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart(content=user_prompt_content)]),
|
||||
ModelResponse(parts=[TextPart(content="I see your image.")]),
|
||||
]
|
||||
|
||||
with patch("chat.agents.local_media_url_processors.update_local_urls") as mock_update:
|
||||
mock_update.return_value = iter([ImageUrl(url="mode-aware-url")])
|
||||
result = update_history_local_urls(conversation, messages)
|
||||
|
||||
assert len(result) == 2
|
||||
mock_update.assert_called_once()
|
||||
|
||||
|
||||
def test_update_local_urls_preserves_other_url_types():
|
||||
"""Test that update_local_urls preserves other URL types unchanged."""
|
||||
conversation = ChatConversationFactory()
|
||||
|
||||
contents = [
|
||||
ImageUrl(url="data:image/png;base64,iVBORw0KG..."), # Already data URL
|
||||
ImageUrl(url="https://example.com/image.jpg"), # HTTPS
|
||||
ImageUrl(url="http://example.com/image.jpg"), # HTTP
|
||||
]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == "data:image/png;base64,iVBORw0KG..."
|
||||
assert result[1].url == "https://example.com/image.jpg"
|
||||
assert result[2].url == "http://example.com/image.jpg"
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_local_urls_stores_updated_urls_mapping(mock_get_file_url):
|
||||
"""Test that update_local_urls stores the mapping of new to old URLs."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "new-mode-aware-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
old_url = f"/media-key/{key}"
|
||||
contents = [ImageUrl(url=old_url)]
|
||||
updated_urls = {}
|
||||
|
||||
list(update_local_urls(conversation, contents, updated_urls))
|
||||
|
||||
assert "new-mode-aware-url" in updated_urls
|
||||
assert updated_urls["new-mode-aware-url"] == old_url
|
||||
|
||||
|
||||
def test_update_local_urls_security_prevents_other_conversation_access():
|
||||
"""Test that security check prevents accessing other conversation's files."""
|
||||
conversation = ChatConversationFactory()
|
||||
other_conversation_key = "other-uuid/attachments/file.jpg"
|
||||
|
||||
# Try to access file from different conversation
|
||||
contents = [ImageUrl(url=f"/media-key/{other_conversation_key}")]
|
||||
|
||||
with patch("chat.agents.local_media_url_processors._get_file_url_for_llm") as mock_get:
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
# URL should not be processed (security check failed)
|
||||
assert result[0].url == f"/media-key/{other_conversation_key}"
|
||||
# _get_file_url_for_llm should NOT be called
|
||||
mock_get.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_get_file_url_is_called_with_correct_arguments(mock_get_file_url):
|
||||
"""Test that _get_file_url_for_llm is called with correct arguments."""
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "processed-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
list(update_local_urls(conversation, contents))
|
||||
|
||||
# Verify the function was called with the S3 key (without /media-key/ prefix)
|
||||
mock_get_file_url.assert_called_once()
|
||||
call_args = mock_get_file_url.call_args
|
||||
assert call_args[0][0] == key # First argument should be the S3 key
|
||||
|
||||
|
||||
# ==================== Tests for FILE_TO_LLM_MODE settings ====================
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_with_presigned_url_mode(mock_policy, settings):
|
||||
"""Test update_local_urls with PRESIGNED_URL mode."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.PRESIGNED_URL
|
||||
conversation = ChatConversationFactory()
|
||||
mock_policy.return_value = "https://s3.example.com/presigned-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].url == "https://s3.example.com/presigned-url"
|
||||
mock_policy.assert_called_once_with(key)
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_temporary_url")
|
||||
def test_update_local_urls_with_backend_temporary_url_mode(mock_temp_url, settings):
|
||||
"""Test update_local_urls with BACKEND_TEMPORARY_URL mode."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_TEMPORARY_URL
|
||||
conversation = ChatConversationFactory()
|
||||
mock_temp_url.return_value = "/api/v1.0/file-stream/temp-token-123/"
|
||||
|
||||
key = f"{conversation.pk}/test.pdf"
|
||||
contents = [DocumentUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].url == "/api/v1.0/file-stream/temp-token-123/"
|
||||
mock_temp_url.assert_called_once_with(key)
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_update_local_urls_with_backend_base64_mode(mock_storage, settings):
|
||||
"""Test update_local_urls with BACKEND_BASE64 mode."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_BASE64
|
||||
conversation = ChatConversationFactory()
|
||||
|
||||
file_content = b"Mock image binary content"
|
||||
mock_file = BytesIO(file_content)
|
||||
mock_storage.return_value.__enter__.return_value = mock_file
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
# Should be a data URL
|
||||
assert result[0].url.startswith("data:")
|
||||
assert "base64" in result[0].url
|
||||
mock_storage.assert_called_once()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_update_local_urls_backend_base64_fallback_on_error(mock_storage, mock_policy, settings):
|
||||
"""Test that update_local_urls falls back to presigned URL when base64 fails."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_BASE64
|
||||
conversation = ChatConversationFactory()
|
||||
mock_storage.side_effect = Exception("S3 connection error")
|
||||
mock_policy.return_value = "https://s3.example.com/fallback-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
# Should fall back to presigned URL
|
||||
assert result[0].url == "https://s3.example.com/fallback-url"
|
||||
mock_policy.assert_called_once_with(key)
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_multiple_files_presigned_mode(mock_policy, settings):
|
||||
"""Test update_local_urls with multiple files in PRESIGNED_URL mode."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.PRESIGNED_URL
|
||||
conversation = ChatConversationFactory()
|
||||
mock_policy.side_effect = [
|
||||
"https://s3.example.com/image1-presigned",
|
||||
"https://s3.example.com/image2-presigned",
|
||||
"https://s3.example.com/document-presigned",
|
||||
]
|
||||
|
||||
key1 = f"{conversation.pk}/image1.jpg"
|
||||
key2 = f"{conversation.pk}/image2.png"
|
||||
key3 = f"{conversation.pk}/document.pdf"
|
||||
|
||||
contents = [
|
||||
ImageUrl(url=f"/media-key/{key1}"),
|
||||
ImageUrl(url=f"/media-key/{key2}"),
|
||||
DocumentUrl(url=f"/media-key/{key3}"),
|
||||
]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == "https://s3.example.com/image1-presigned"
|
||||
assert result[1].url == "https://s3.example.com/image2-presigned"
|
||||
assert result[2].url == "https://s3.example.com/document-presigned"
|
||||
assert mock_policy.call_count == 3
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_temporary_url")
|
||||
def test_update_local_urls_multiple_files_temporary_url_mode(mock_temp_url, settings):
|
||||
"""Test update_local_urls with multiple files in BACKEND_TEMPORARY_URL mode."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_TEMPORARY_URL
|
||||
conversation = ChatConversationFactory()
|
||||
mock_temp_url.side_effect = [
|
||||
"/api/v1.0/file-stream/token1/",
|
||||
"/api/v1.0/file-stream/token2/",
|
||||
"/api/v1.0/file-stream/token3/",
|
||||
]
|
||||
|
||||
key1 = f"{conversation.pk}/image1.jpg"
|
||||
key2 = f"{conversation.pk}/image2.png"
|
||||
key3 = f"{conversation.pk}/document.pdf"
|
||||
|
||||
contents = [
|
||||
ImageUrl(url=f"/media-key/{key1}"),
|
||||
ImageUrl(url=f"/media-key/{key2}"),
|
||||
DocumentUrl(url=f"/media-key/{key3}"),
|
||||
]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 3
|
||||
assert result[0].url == "/api/v1.0/file-stream/token1/"
|
||||
assert result[1].url == "/api/v1.0/file-stream/token2/"
|
||||
assert result[2].url == "/api/v1.0/file-stream/token3/"
|
||||
assert mock_temp_url.call_count == 3
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_presigned_mode_preserves_mapping(mock_policy, settings):
|
||||
"""Test that PRESIGNED_URL mode correctly stores updated URLs mapping."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.PRESIGNED_URL
|
||||
conversation = ChatConversationFactory()
|
||||
presigned_url = "https://s3.example.com/presigned-123"
|
||||
mock_policy.return_value = presigned_url
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
original_url = f"/media-key/{key}"
|
||||
contents = [ImageUrl(url=original_url)]
|
||||
updated_urls = {}
|
||||
|
||||
list(update_local_urls(conversation, contents, updated_urls))
|
||||
|
||||
assert presigned_url in updated_urls
|
||||
assert updated_urls[presigned_url] == original_url
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_temporary_url")
|
||||
def test_update_local_urls_temporary_url_mode_preserves_mapping(mock_temp_url, settings):
|
||||
"""Test that BACKEND_TEMPORARY_URL mode correctly stores updated URLs mapping."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_TEMPORARY_URL
|
||||
conversation = ChatConversationFactory()
|
||||
temp_url = "/api/v1.0/file-stream/temp-abc-def/"
|
||||
mock_temp_url.return_value = temp_url
|
||||
|
||||
key = f"{conversation.pk}/test.pdf"
|
||||
original_url = f"/media-key/{key}"
|
||||
contents = [DocumentUrl(url=original_url)]
|
||||
updated_urls = {}
|
||||
|
||||
list(update_local_urls(conversation, contents, updated_urls))
|
||||
|
||||
assert temp_url in updated_urls
|
||||
assert updated_urls[temp_url] == original_url
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_update_local_urls_base64_mode_preserves_mapping(mock_storage, settings):
|
||||
"""Test that BACKEND_BASE64 mode correctly stores updated URLs mapping."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_BASE64
|
||||
conversation = ChatConversationFactory()
|
||||
|
||||
file_content = b"test image content"
|
||||
mock_file = BytesIO(file_content)
|
||||
mock_storage.return_value.__enter__.return_value = mock_file
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
original_url = f"/media-key/{key}"
|
||||
contents = [ImageUrl(url=original_url)]
|
||||
updated_urls = {}
|
||||
|
||||
result = list(update_local_urls(conversation, contents, updated_urls))
|
||||
|
||||
# Verify mapping was stored
|
||||
assert len(updated_urls) == 1
|
||||
# The data URL should be the key in the mapping
|
||||
data_url = result[0].url
|
||||
assert data_url in updated_urls
|
||||
assert updated_urls[data_url] == original_url
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_presigned_mode_security_check(mock_policy, settings):
|
||||
"""Test that PRESIGNED_URL mode respects security checks."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.PRESIGNED_URL
|
||||
conversation = ChatConversationFactory()
|
||||
other_key = "other-conversation-id/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{other_key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
# URL should remain unchanged
|
||||
assert result[0].url == f"/media-key/{other_key}"
|
||||
# generate_retrieve_policy should not be called
|
||||
mock_policy.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_temporary_url")
|
||||
def test_update_local_urls_temporary_mode_security_check(mock_temp_url, settings):
|
||||
"""Test that BACKEND_TEMPORARY_URL mode respects security checks."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_TEMPORARY_URL
|
||||
conversation = ChatConversationFactory()
|
||||
other_key = "other-conversation-id/test.pdf"
|
||||
contents = [DocumentUrl(url=f"/media-key/{other_key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
# URL should remain unchanged
|
||||
assert result[0].url == f"/media-key/{other_key}"
|
||||
# generate_temporary_url should not be called
|
||||
mock_temp_url.assert_not_called()
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_update_local_urls_base64_mode_security_check(mock_storage, settings):
|
||||
"""Test that BACKEND_BASE64 mode respects security checks."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_BASE64
|
||||
conversation = ChatConversationFactory()
|
||||
other_key = "other-conversation-id/test.jpg"
|
||||
contents = [ImageUrl(url=f"/media-key/{other_key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
# URL should remain unchanged
|
||||
assert result[0].url == f"/media-key/{other_key}"
|
||||
# Storage should not be opened
|
||||
mock_storage.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_to_llm_mode",
|
||||
[
|
||||
FileToLLMMode.PRESIGNED_URL,
|
||||
FileToLLMMode.BACKEND_TEMPORARY_URL,
|
||||
FileToLLMMode.BACKEND_BASE64,
|
||||
],
|
||||
)
|
||||
@patch("chat.agents.local_media_url_processors._get_file_url_for_llm")
|
||||
def test_update_local_urls_all_modes_with_external_urls(
|
||||
mock_get_file_url, file_to_llm_mode, settings
|
||||
):
|
||||
"""Test that all modes preserve external URLs unchanged."""
|
||||
settings.FILE_TO_LLM_MODE = file_to_llm_mode
|
||||
conversation = ChatConversationFactory()
|
||||
mock_get_file_url.return_value = "processed-url"
|
||||
|
||||
key = f"{conversation.pk}/test.jpg"
|
||||
external_urls = [
|
||||
"https://example.com/image.jpg",
|
||||
"http://another.com/image.png",
|
||||
"data:image/png;base64,iVBORw0KG...",
|
||||
]
|
||||
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")] + [ImageUrl(url=url) for url in external_urls]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 4
|
||||
assert result[0].url == "processed-url"
|
||||
for i, external_url in enumerate(external_urls, start=1):
|
||||
assert result[i].url == external_url
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.default_storage.open")
|
||||
def test_update_local_urls_base64_mode_with_different_file_types(mock_storage, settings):
|
||||
"""Test BACKEND_BASE64 mode with different file MIME types."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.BACKEND_BASE64
|
||||
conversation = ChatConversationFactory()
|
||||
|
||||
# Test with different file types
|
||||
test_cases = [
|
||||
("test.jpg", b"JPEG binary"),
|
||||
("test.png", b"PNG binary"),
|
||||
("test.pdf", b"PDF binary"),
|
||||
("test.txt", b"Text content"),
|
||||
]
|
||||
|
||||
for filename, content in test_cases:
|
||||
mock_file = BytesIO(content)
|
||||
mock_storage.return_value.__enter__.return_value = mock_file
|
||||
|
||||
key = f"{conversation.pk}/{filename}"
|
||||
contents = [ImageUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
# Should be a data URL
|
||||
assert result[0].url.startswith("data:")
|
||||
assert "base64" in result[0].url
|
||||
|
||||
|
||||
@patch("chat.agents.local_media_url_processors.generate_retrieve_policy")
|
||||
def test_update_local_urls_presigned_mode_with_special_characters_in_key(mock_policy, settings):
|
||||
"""Test PRESIGNED_URL mode handles keys with special characters."""
|
||||
settings.FILE_TO_LLM_MODE = FileToLLMMode.PRESIGNED_URL
|
||||
conversation = ChatConversationFactory()
|
||||
mock_policy.return_value = "https://s3.example.com/presigned"
|
||||
|
||||
# Key with special characters (should be handled properly by S3)
|
||||
key = f"{conversation.pk}/attachments/file (1).pdf"
|
||||
contents = [DocumentUrl(url=f"/media-key/{key}")]
|
||||
|
||||
result = list(update_local_urls(conversation, contents))
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0].url == "https://s3.example.com/presigned"
|
||||
mock_policy.assert_called_once_with(key)
|
||||
|
||||
@@ -9,7 +9,7 @@ from django.test import override_settings
|
||||
|
||||
import pytest
|
||||
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.enums import AttachmentStatus, FileUploadMode
|
||||
|
||||
from chat import factories, models
|
||||
from chat.tests.conftest import PIXEL_PNG
|
||||
@@ -266,3 +266,40 @@ def test_upload_ended_fix_extension(api_client, name, content, _extension, conte
|
||||
assert attachment.content_type == content_type # updated
|
||||
assert attachment.file_name == name # updated
|
||||
assert attachment.size == len(content) # updated
|
||||
|
||||
|
||||
def test_attachment_create_presigned_url_mode_returns_policy(api_client, settings):
|
||||
"""Test that presigned_url mode returns policy field."""
|
||||
settings.FILE_UPLOAD_MODE = FileUploadMode.PRESIGNED_URL
|
||||
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
response = api_client.post(
|
||||
url,
|
||||
{"file_name": "test.pdf", "size": 1000, "content_type": "application/pdf"},
|
||||
format="json",
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
data = response.json()
|
||||
assert data["policy"] is not None
|
||||
assert "s3" in data["policy"].lower() or "minio" in data["policy"].lower()
|
||||
|
||||
|
||||
def test_attachment_create_backend_to_s3_mode_no_policy(api_client, settings):
|
||||
"""Test that backend_to_s3 mode does not return policy."""
|
||||
settings.FILE_UPLOAD_MODE = FileUploadMode.BACKEND_TO_S3
|
||||
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/"
|
||||
response = api_client.post(
|
||||
url, {"file_name": "test.pdf", "content_type": "application/pdf"}, format="json"
|
||||
)
|
||||
|
||||
assert response.status_code == 201
|
||||
data = response.json()
|
||||
assert data["policy"] is None
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Tests for file upload mode `backend_to_s3`."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest import mock
|
||||
|
||||
from django.core.files.storage import default_storage
|
||||
|
||||
import pytest
|
||||
|
||||
from chat import factories, models
|
||||
from chat.tests.conftest import PIXEL_PNG
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
def test_backend_upload_anonymous_forbidden(api_client):
|
||||
"""Anonymous users should not be able to use backend upload."""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/backend-upload/"
|
||||
|
||||
response = api_client.post(url, {}, format="multipart")
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
def test_backend_upload_not_owner_forbidden(api_client):
|
||||
"""Users who don't own the conversation cannot upload via backend."""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
user = factories.UserFactory()
|
||||
api_client.force_login(user)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/backend-upload/"
|
||||
file_obj = BytesIO(PIXEL_PNG)
|
||||
file_obj.name = "test.png"
|
||||
|
||||
response = api_client.post(url, {"file": file_obj, "file_name": "test.png"}, format="multipart")
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@mock.patch("chat.views.malware_detection.analyse_file")
|
||||
def test_backend_upload_success(mock_malware, api_client):
|
||||
"""Test successful backend file upload."""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/backend-upload/"
|
||||
file_obj = BytesIO(PIXEL_PNG)
|
||||
file_obj.name = "test.png"
|
||||
|
||||
response = api_client.post(url, {"file": file_obj, "file_name": "test.png"}, format="multipart")
|
||||
|
||||
assert response.status_code == 201
|
||||
data = response.json()
|
||||
|
||||
# Verify response structure
|
||||
assert "id" in data
|
||||
assert "key" in data
|
||||
assert data["file_name"] == "test.png"
|
||||
assert data["size"] == len(PIXEL_PNG)
|
||||
assert data["content_type"] == "image/png"
|
||||
|
||||
# Verify attachment was created
|
||||
attachment = models.ChatConversationAttachment.objects.get(pk=data["id"])
|
||||
assert attachment.conversation == conversation
|
||||
assert attachment.uploaded_by == conversation.owner
|
||||
assert attachment.file_name == "test.png"
|
||||
assert attachment.size == len(PIXEL_PNG)
|
||||
|
||||
# Verify malware detection was called
|
||||
mock_malware.assert_called_once()
|
||||
|
||||
|
||||
def test_backend_upload_missing_file(api_client):
|
||||
"""Test that backend upload fails without file field."""
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/backend-upload/"
|
||||
|
||||
response = api_client.post(
|
||||
url,
|
||||
{"file_name": "test.png"}, # No file field
|
||||
format="multipart",
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "file" in response.json()
|
||||
|
||||
|
||||
@mock.patch("chat.views.malware_detection.analyse_file")
|
||||
def test_backend_upload_creates_s3_file(_mock_malware, api_client):
|
||||
"""Test that backend upload creates file in S3."""
|
||||
|
||||
conversation = factories.ChatConversationFactory()
|
||||
api_client.force_login(conversation.owner)
|
||||
|
||||
url = f"/api/v1.0/chats/{conversation.pk!s}/attachments/backend-upload/"
|
||||
file_obj = BytesIO(PIXEL_PNG)
|
||||
file_obj.name = "test.png"
|
||||
|
||||
response = api_client.post(url, {"file": file_obj, "file_name": "test.png"}, format="multipart")
|
||||
|
||||
assert response.status_code == 201
|
||||
data = response.json()
|
||||
key = data["key"]
|
||||
|
||||
# Verify file exists in S3
|
||||
assert default_storage.exists(key)
|
||||
|
||||
# Verify file content
|
||||
with default_storage.open(key, "rb") as f:
|
||||
content = f.read()
|
||||
assert content == PIXEL_PNG
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Tests for the file stream endpoint."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest import mock
|
||||
|
||||
from django.core.cache import cache
|
||||
|
||||
|
||||
def test_file_stream_invalid_key(api_client):
|
||||
"""Test that invalid temporary keys return 404."""
|
||||
cache.clear()
|
||||
url = "/api/v1.0/file-stream/invalid-key/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 404
|
||||
error = response.json()["detail"].lower()
|
||||
assert "expired" in error or "invalid" in error
|
||||
|
||||
|
||||
def test_file_stream_expired_key(api_client):
|
||||
"""Test that expired keys return 404."""
|
||||
cache.clear()
|
||||
# Create a key that's already expired
|
||||
cache.set("file_access:expired-key", "path/to/file.pdf", timeout=0)
|
||||
|
||||
url = "/api/v1.0/file-stream/expired-key/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
@mock.patch("chat.views.magic.Magic")
|
||||
@mock.patch("chat.views.default_storage.open")
|
||||
def test_file_stream_valid_key_streams_file(mock_storage_open, mock_magic, api_client):
|
||||
"""Test that valid temporary keys stream file content."""
|
||||
cache.clear()
|
||||
|
||||
# Create a valid temporary key
|
||||
temporary_key = "test-valid-key"
|
||||
s3_key = "test/path/file.pdf"
|
||||
cache.set(f"file_access:{temporary_key}", s3_key, timeout=300)
|
||||
|
||||
# Mock storage.open to return file content
|
||||
file_mock = BytesIO(b"PDF content here")
|
||||
mock_storage_open.return_value = file_mock
|
||||
|
||||
# Mock magic detector
|
||||
mock_magic_instance = mock.MagicMock()
|
||||
mock_magic_instance.from_buffer.return_value = "application/pdf"
|
||||
mock_magic.return_value = mock_magic_instance
|
||||
|
||||
url = f"/api/v1.0/file-stream/{temporary_key}/"
|
||||
response = api_client.get(url)
|
||||
|
||||
assert response.status_code == 200
|
||||
@@ -5,6 +5,7 @@ import os
|
||||
from uuid import uuid4
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
from django.core.files.storage import default_storage
|
||||
from django.http import Http404, StreamingHttpResponse
|
||||
|
||||
@@ -15,12 +16,14 @@ from lasuite.malware_detection import malware_detection
|
||||
from rest_framework import decorators, filters, mixins, permissions, status, viewsets
|
||||
from rest_framework.exceptions import MethodNotAllowed, PermissionDenied, ValidationError
|
||||
from rest_framework.response import Response
|
||||
from rest_framework.throttling import ScopedRateThrottle
|
||||
from rest_framework.views import APIView
|
||||
|
||||
from core.api.viewsets import Pagination, SerializerPerActionMixin
|
||||
from core.file_upload import enums
|
||||
from core.file_upload.enums import AttachmentStatus
|
||||
from core.file_upload.mixins import AttachmentMixin
|
||||
from core.file_upload.serializers import FileUploadSerializer
|
||||
from core.filters import remove_accents
|
||||
|
||||
from activation_codes.permissions import IsActivatedUser
|
||||
@@ -434,3 +437,173 @@ class ChatConversationAttachmentViewSet(
|
||||
)
|
||||
|
||||
return Response(serializer.data, status=status.HTTP_200_OK)
|
||||
|
||||
@decorators.action(
|
||||
detail=False,
|
||||
methods=["post"],
|
||||
url_path="backend-upload",
|
||||
url_name="backend-upload",
|
||||
)
|
||||
def backend_upload_attachment(self, request, *args, **kwargs):
|
||||
"""
|
||||
Handle backend file upload for backend_to_s3 mode.
|
||||
|
||||
This endpoint is used when FILE_UPLOAD_MODE is set to backend_to_s3.
|
||||
The frontend sends the file directly to this endpoint,
|
||||
and the backend stores it on S3 and initiates malware detection.
|
||||
|
||||
The attachment lifecycle:
|
||||
1. Frontend sends file via this endpoint
|
||||
2. Backend stores file on S3
|
||||
3. Backend detects MIME type and file size
|
||||
4. Backend initiates malware detection
|
||||
5. After detection, attachment status becomes READY or SUSPICIOUS
|
||||
"""
|
||||
# pylint: disable=too-many-locals
|
||||
# Verify the user owns the conversation
|
||||
conversation_id = self.kwargs["conversation_pk"]
|
||||
if not models.ChatConversation.objects.filter(
|
||||
pk=conversation_id,
|
||||
owner=request.user,
|
||||
).exists():
|
||||
raise Http404
|
||||
|
||||
serializer = FileUploadSerializer(data=request.data)
|
||||
serializer.is_valid(raise_exception=True)
|
||||
|
||||
file_obj = serializer.validated_data["file"]
|
||||
file_name = serializer.validated_data["file_name"]
|
||||
|
||||
# Generate unique file ID and storage key
|
||||
file_id = uuid4()
|
||||
extension = file_name.rpartition(".")[-1] if "." in file_name else None
|
||||
ext_suffix = f".{extension}" if extension else ""
|
||||
key = f"{conversation_id}/{AttachmentMixin.ATTACHMENTS_FOLDER}/{file_id}{ext_suffix}"
|
||||
|
||||
# Store file on S3
|
||||
try:
|
||||
stored_path = default_storage.save(key, file_obj)
|
||||
logger.info("File uploaded to S3: %s", stored_path)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
logger.exception("Failed to upload file to S3 for conversation %s", conversation_id)
|
||||
return Response(
|
||||
{"detail": "Failed to upload file to storage"},
|
||||
status=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
)
|
||||
|
||||
# Detect MIME type
|
||||
mime_detector = magic.Magic(mime=True)
|
||||
with default_storage.open(key, "rb") as file:
|
||||
mimetype = mime_detector.from_buffer(file.read(2048))
|
||||
file_size = file.size
|
||||
|
||||
# Create attachment record with ANALYZING status
|
||||
attachment = models.ChatConversationAttachment.objects.create(
|
||||
conversation_id=conversation_id,
|
||||
uploaded_by=request.user,
|
||||
upload_state=AttachmentStatus.ANALYZING,
|
||||
key=key,
|
||||
file_name=file_name,
|
||||
content_type=mimetype,
|
||||
size=file_size,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Created attachment %s for conversation %s, starting malware detection",
|
||||
attachment.pk,
|
||||
conversation_id,
|
||||
)
|
||||
|
||||
# Start malware detection (will update status to READY or SUSPICIOUS via callbacks)
|
||||
malware_detection.analyse_file(
|
||||
key,
|
||||
safe_callback="chat.malware_detection.conversation_safe_attachment_callback",
|
||||
unknown_callback="chat.malware_detection.unknown_attachment_callback",
|
||||
unsafe_callback="chat.malware_detection.conversation_unsafe_attachment_callback",
|
||||
conversation_id=conversation_id,
|
||||
)
|
||||
|
||||
# Track upload event
|
||||
if settings.POSTHOG_KEY:
|
||||
posthog.capture(
|
||||
"item_uploaded_backend",
|
||||
distinct_id=str(request.user.pk),
|
||||
properties={
|
||||
"id": attachment.pk,
|
||||
"file_name": attachment.file_name,
|
||||
"size": attachment.size,
|
||||
"mimetype": attachment.content_type,
|
||||
"mode": settings.FILE_UPLOAD_MODE,
|
||||
},
|
||||
)
|
||||
|
||||
serializer = self.get_serializer(attachment)
|
||||
return Response(serializer.data, status=status.HTTP_201_CREATED)
|
||||
|
||||
|
||||
class FileStreamView(APIView):
|
||||
"""
|
||||
Stream file content for temporary access URLs.
|
||||
|
||||
This view is used by LLMs to access file content when they cannot directly
|
||||
access S3. A temporary key is stored in cache and validated before serving
|
||||
the file.
|
||||
|
||||
Security:
|
||||
- Temporary key expires after FILE_BACKEND_TEMPORARY_URL_EXPIRATION seconds
|
||||
(default: 180 seconds / 3 minutes)
|
||||
- No authentication required (key is single-use temporary token)
|
||||
- Key is generated using secure random tokens
|
||||
"""
|
||||
|
||||
permission_classes = [] # No authentication needed for temporary keys
|
||||
throttle_classes = [ScopedRateThrottle]
|
||||
throttle_scope = "file-stream"
|
||||
|
||||
def get(self, request, temporary_key):
|
||||
"""
|
||||
Stream file content using a temporary access key.
|
||||
|
||||
Args:
|
||||
temporary_key: The temporary key generated by generate_temporary_url()
|
||||
|
||||
Returns:
|
||||
StreamingHttpResponse with file content
|
||||
"""
|
||||
# Retrieve the S3 key from cache using the temporary key
|
||||
cache_key = f"file_access:{temporary_key}"
|
||||
s3_key = cache.get(cache_key)
|
||||
|
||||
if not s3_key:
|
||||
logger.warning("Temporary file access key not found or expired: %s", temporary_key)
|
||||
raise Http404("File access key expired or invalid")
|
||||
|
||||
# Delete the key from cache to prevent reuse
|
||||
cache.delete(cache_key)
|
||||
|
||||
logger.info("Serving file via temporary key: %s", s3_key)
|
||||
|
||||
try:
|
||||
# Open the file from S3
|
||||
file_obj = default_storage.open(s3_key, "rb")
|
||||
|
||||
# Detect MIME type for proper content-type header
|
||||
mime_detector = magic.Magic(mime=True)
|
||||
file_content = file_obj.read(2048)
|
||||
file_obj.seek(0)
|
||||
content_type = mime_detector.from_buffer(file_content)
|
||||
|
||||
# Extract filename from S3 key (last part after /)
|
||||
filename = s3_key.split("/")[-1]
|
||||
|
||||
# Stream the file content
|
||||
response = StreamingHttpResponse(
|
||||
file_obj,
|
||||
content_type=content_type,
|
||||
)
|
||||
response["Content-Disposition"] = f'inline; filename="{filename}"'
|
||||
return response
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("Failed to serve file via temporary key: %s", temporary_key)
|
||||
raise Http404("Failed to retrieve file") from exc
|
||||
|
||||
@@ -230,6 +230,27 @@ class Base(BraveSettings, Configuration):
|
||||
environ_name="ATTACHMENT_MAX_SIZE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
FILE_UPLOAD_MODE = values.Value(
|
||||
"presigned_url",
|
||||
environ_name="FILE_UPLOAD_MODE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
FILE_TO_LLM_MODE = values.Value(
|
||||
"presigned_url",
|
||||
environ_name="FILE_TO_LLM_MODE",
|
||||
environ_prefix=None,
|
||||
)
|
||||
FILE_BACKEND_URL = values.Value(
|
||||
"",
|
||||
environ_name="FILE_BACKEND_URL",
|
||||
environ_prefix=None,
|
||||
)
|
||||
FILE_BACKEND_TEMPORARY_URL_EXPIRATION = values.IntegerValue(
|
||||
180,
|
||||
environ_name="FILE_BACKEND_TEMPORARY_URL_EXPIRATION",
|
||||
environ_prefix=None,
|
||||
)
|
||||
|
||||
MALWARE_DETECTION = {
|
||||
"BACKEND": values.Value(
|
||||
"lasuite.malware_detection.backends.dummy.DummyBackend",
|
||||
@@ -395,6 +416,11 @@ class Base(BraveSettings, Configuration):
|
||||
environ_name="API_USERS_LIST_THROTTLE_RATE_BURST",
|
||||
environ_prefix=None,
|
||||
),
|
||||
"file-stream": values.Value(
|
||||
default="60/minute",
|
||||
environ_name="API_FILE_STREAM_THROTTLE_RATE",
|
||||
environ_prefix=None,
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -1046,6 +1072,13 @@ USER QUESTION:
|
||||
"OIDC_ALLOW_DUPLICATE_EMAILS cannot be set to True simultaneously. "
|
||||
)
|
||||
|
||||
# File access configuration validation
|
||||
if cls.FILE_TO_LLM_MODE == "backend_temporary_url" and not cls.FILE_BACKEND_URL:
|
||||
raise ValueError(
|
||||
"FILE_TO_LLM_MODE is set to 'backend_temporary_url' but FILE_BACKEND_URL is empty. "
|
||||
"Please set FILE_BACKEND_URL to a valid URL for backend temporary file access."
|
||||
)
|
||||
|
||||
# Langfuse initialization
|
||||
if cls.LANGFUSE_ENABLED:
|
||||
if not cls.LANGFUSE_MEDIA_UPLOAD_ENABLED:
|
||||
|
||||
@@ -211,6 +211,7 @@ class ConfigView(drf.views.APIView):
|
||||
"LANGUAGE_CODE",
|
||||
"SENTRY_DSN",
|
||||
"FEATURE_FLAGS",
|
||||
"FILE_UPLOAD_MODE",
|
||||
]
|
||||
dict_settings = {}
|
||||
for setting in array_settings:
|
||||
|
||||
@@ -20,3 +20,28 @@ class AttachmentStatus(StrEnum):
|
||||
def choices(cls):
|
||||
"""Return a list of tuples for each enum member."""
|
||||
return [(member.value, member.name) for member in cls]
|
||||
|
||||
|
||||
class FileUploadMode(StrEnum):
|
||||
"""Defines the possible modes for file upload (from frontend) handling."""
|
||||
|
||||
PRESIGNED_URL = "presigned_url"
|
||||
BACKEND_TO_S3 = "backend_to_s3"
|
||||
|
||||
@classmethod
|
||||
def choices(cls):
|
||||
"""Return a list of tuples for each enum member."""
|
||||
return [(member.value, member.name) for member in cls]
|
||||
|
||||
|
||||
class FileToLLMMode(StrEnum):
|
||||
"""Defines the possible modes to send file to the LLM."""
|
||||
|
||||
PRESIGNED_URL = "presigned_url"
|
||||
BACKEND_BASE64 = "backend_base64"
|
||||
BACKEND_TEMPORARY_URL = "backend_temporary_url"
|
||||
|
||||
@classmethod
|
||||
def choices(cls):
|
||||
"""Return a list of tuples for each enum member."""
|
||||
return [(member.value, member.name) for member in cls]
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
"""Tests for generate_temporary_url utility function."""
|
||||
|
||||
from django.conf import settings
|
||||
from django.core.cache import cache
|
||||
|
||||
import pytest
|
||||
|
||||
from chat.agents.local_media_url_processors import generate_temporary_url
|
||||
|
||||
pytestmark = pytest.mark.django_db
|
||||
|
||||
|
||||
def test_generate_temporary_url_returns_string():
|
||||
"""Test that generate_temporary_url returns a valid backend streaming URL."""
|
||||
cache.clear()
|
||||
url = generate_temporary_url("test/file.pdf")
|
||||
|
||||
assert isinstance(url, str)
|
||||
assert url.startswith(settings.FILE_BACKEND_URL + "/api/v1.0/file-stream/")
|
||||
assert url.endswith("/")
|
||||
|
||||
|
||||
def test_generate_temporary_url_creates_cache_entry():
|
||||
"""Test that a cache entry is created with correct mapping."""
|
||||
cache.clear()
|
||||
s3_key = "conversation-id/attachments/file-uuid.pdf"
|
||||
url = generate_temporary_url(s3_key)
|
||||
|
||||
# Extract temporary key from URL
|
||||
temporary_key = url.split("/file-stream/")[1].rstrip("/")
|
||||
|
||||
# Verify cache entry
|
||||
cache_key = f"file_access:{temporary_key}"
|
||||
cached_value = cache.get(cache_key)
|
||||
assert cached_value == s3_key
|
||||
|
||||
|
||||
def test_generate_temporary_url_unique_tokens():
|
||||
"""Test that different S3 keys produce different temporary tokens."""
|
||||
cache.clear()
|
||||
url1 = generate_temporary_url("file1.pdf")
|
||||
url2 = generate_temporary_url("file2.pdf")
|
||||
|
||||
assert url1 != url2
|
||||
|
||||
key1 = url1.split("/file-stream/")[1].rstrip("/")
|
||||
key2 = url2.split("/file-stream/")[1].rstrip("/")
|
||||
|
||||
assert cache.get(f"file_access:{key1}") == "file1.pdf"
|
||||
assert cache.get(f"file_access:{key2}") == "file2.pdf"
|
||||
|
||||
|
||||
def test_generate_temporary_url_token_is_url_safe():
|
||||
"""Test that generated tokens contain only URL-safe characters."""
|
||||
cache.clear()
|
||||
url = generate_temporary_url("test.pdf")
|
||||
|
||||
temporary_key = url.split("/file-stream/")[1].rstrip("/")
|
||||
|
||||
# Token should only contain alphanumeric, dash, and underscore
|
||||
assert all(c.isalnum() or c in "-_" for c in temporary_key)
|
||||
|
||||
|
||||
def test_generate_temporary_url_token_sufficient_entropy():
|
||||
"""Test that generated tokens have sufficient entropy."""
|
||||
cache.clear()
|
||||
url = generate_temporary_url("test.pdf")
|
||||
|
||||
temporary_key = url.split("/file-stream/")[1].rstrip("/")
|
||||
|
||||
# Token should be reasonably long
|
||||
assert len(temporary_key) >= 32
|
||||
|
||||
|
||||
def test_generate_temporary_url_no_sensitive_data_in_url():
|
||||
"""Test that temporary URLs don't contain S3 key information."""
|
||||
cache.clear()
|
||||
|
||||
s3_key = "secret/conversation-123/attachments/file.pdf"
|
||||
url = generate_temporary_url(s3_key)
|
||||
|
||||
# URL should not contain the actual S3 key
|
||||
assert "secret" not in url
|
||||
assert "conversation-123" not in url
|
||||
assert "file.pdf" not in url
|
||||
# Only the endpoint and random token
|
||||
assert "/api/v1.0/file-stream/" in url
|
||||
|
||||
|
||||
def test_generate_temporary_url_various_key_formats():
|
||||
"""Test generate_temporary_url with various S3 key formats."""
|
||||
cache.clear()
|
||||
|
||||
test_keys = [
|
||||
"simple/key.pdf",
|
||||
"conversation-123/attachments/file-uuid.pdf",
|
||||
"nested/folder/structure/file.jpg",
|
||||
"file_with_special-chars_123.png",
|
||||
]
|
||||
|
||||
urls = []
|
||||
for key in test_keys:
|
||||
url = generate_temporary_url(key)
|
||||
urls.append(url)
|
||||
|
||||
temporary_key = url.split("/file-stream/")[1].rstrip("/")
|
||||
assert cache.get(f"file_access:{temporary_key}") == key
|
||||
|
||||
# All URLs should be different
|
||||
assert len(set(urls)) == len(urls)
|
||||
@@ -47,6 +47,7 @@ def test_api_config(is_authenticated):
|
||||
"CRISP_WEBSITE_ID": "123",
|
||||
"ENVIRONMENT": "test",
|
||||
"FEATURE_FLAGS": {"document-upload": "enabled", "web-search": "enabled"},
|
||||
"FILE_UPLOAD_MODE": "presigned_url",
|
||||
"FRONTEND_CSS_URL": "http://testcss/",
|
||||
"FRONTEND_HOMEPAGE_FEATURE_ENABLED": True,
|
||||
"FRONTEND_THEME": "test-theme",
|
||||
@@ -189,6 +190,7 @@ async def test_api_config_async(is_authenticated):
|
||||
"CRISP_WEBSITE_ID": "123",
|
||||
"ENVIRONMENT": "test",
|
||||
"FEATURE_FLAGS": {"document-upload": "enabled", "web-search": "enabled"},
|
||||
"FILE_UPLOAD_MODE": "presigned_url",
|
||||
"FRONTEND_CSS_URL": "http://testcss/",
|
||||
"FRONTEND_HOMEPAGE_FEATURE_ENABLED": True,
|
||||
"FRONTEND_THEME": "test-theme",
|
||||
|
||||
@@ -1,15 +1,21 @@
|
||||
"""URL configuration for the core app."""
|
||||
|
||||
from django.conf import settings
|
||||
from django.urls import include, path
|
||||
from django.urls import include, path, re_path
|
||||
|
||||
from lasuite.oidc_login.urls import urlpatterns as oidc_urls
|
||||
from rest_framework.routers import DefaultRouter
|
||||
|
||||
from core.api import viewsets
|
||||
from core.file_upload.enums import FileToLLMMode
|
||||
|
||||
from activation_codes import viewsets as activation_viewsets
|
||||
from chat.views import ChatConversationAttachmentViewSet, ChatViewSet, LLMConfigurationView
|
||||
from chat.views import (
|
||||
ChatConversationAttachmentViewSet,
|
||||
ChatViewSet,
|
||||
FileStreamView,
|
||||
LLMConfigurationView,
|
||||
)
|
||||
|
||||
# - Main endpoints
|
||||
router = DefaultRouter()
|
||||
@@ -37,6 +43,21 @@ urlpatterns = [
|
||||
include(conversation_router.urls),
|
||||
),
|
||||
]
|
||||
+ (
|
||||
# Only allow file stream URL when configured for backend temporary URL mode
|
||||
[
|
||||
re_path(
|
||||
r"^file-stream/(?P<temporary_key>[a-zA-Z0-9_-]+)/$",
|
||||
FileStreamView.as_view(),
|
||||
name="file-stream",
|
||||
),
|
||||
]
|
||||
if (
|
||||
settings.FILE_TO_LLM_MODE == FileToLLMMode.BACKEND_TEMPORARY_URL
|
||||
or settings.ENVIRONMENT == "test"
|
||||
)
|
||||
else []
|
||||
)
|
||||
),
|
||||
),
|
||||
path(f"api/{settings.API_VERSION}/config/", viewsets.ConfigView.as_view()),
|
||||
|
||||
@@ -34,6 +34,7 @@ export interface ConfigResponse {
|
||||
MEDIA_BASE_URL?: string;
|
||||
POSTHOG_KEY?: PostHogConf;
|
||||
SENTRY_DSN?: string;
|
||||
FILE_UPLOAD_MODE?: string;
|
||||
theme_customization?: ThemeCustomization;
|
||||
chat_upload_accept?: string;
|
||||
}
|
||||
|
||||
@@ -1,14 +1,19 @@
|
||||
import { useCallback } from 'react';
|
||||
|
||||
import { fetchAPI } from '@/api';
|
||||
import { baseApiUrl, fetchAPI, getCSRFToken } from '@/api';
|
||||
import { useConfig } from '@/core';
|
||||
|
||||
import { useCreateConversationAttachment } from '../api';
|
||||
|
||||
interface BackendUploadResponse {
|
||||
key: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Upload a file, using XHR so we can report on progress through a handler.
|
||||
* @param url The pre-signed URL to PUT the file to.
|
||||
* @param file The raw file to upload as the request body.
|
||||
* @param progressHandler A handler that receives progress updates as a single integer `0 <= x <= 100`.
|
||||
* @param progressHandler A handler that receives progress updates as a single integer `0 <= x <= 100`.
|
||||
*/
|
||||
export const uploadFileToServer = (
|
||||
url: string,
|
||||
@@ -47,6 +52,71 @@ export const uploadFileToServer = (
|
||||
xhr.send(file);
|
||||
});
|
||||
|
||||
/**
|
||||
* Upload a file to the backend (for backend_base64 and backend_temporary_url modes).
|
||||
* Uses XHR to track upload progress while respecting the project's API patterns.
|
||||
* @param conversationId The ID of the conversation.
|
||||
* @param file The file to upload.
|
||||
* @param progressHandler A handler that receives progress updates.
|
||||
*/
|
||||
export const uploadFileToBackend = (
|
||||
conversationId: string,
|
||||
file: File,
|
||||
progressHandler: (progress: number) => void,
|
||||
): Promise<BackendUploadResponse> =>
|
||||
new Promise((resolve, reject) => {
|
||||
const formData = new FormData();
|
||||
formData.append('file', file);
|
||||
formData.append('file_name', file.name);
|
||||
|
||||
const xhr = new XMLHttpRequest();
|
||||
const csrfToken = getCSRFToken();
|
||||
|
||||
xhr.addEventListener('error', reject);
|
||||
xhr.addEventListener('abort', reject);
|
||||
|
||||
xhr.addEventListener('readystatechange', () => {
|
||||
if (xhr.readyState === 4) {
|
||||
if (xhr.status === 201) {
|
||||
progressHandler(100);
|
||||
try {
|
||||
const response = JSON.parse(
|
||||
xhr.responseText,
|
||||
) as BackendUploadResponse;
|
||||
return resolve(response);
|
||||
} catch {
|
||||
return reject(new Error('Failed to parse server response'));
|
||||
}
|
||||
}
|
||||
reject(
|
||||
new Error(
|
||||
`Failed to upload file to backend: ${xhr.status} ${xhr.statusText}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
xhr.upload.addEventListener('progress', (progressEvent) => {
|
||||
if (progressEvent.lengthComputable) {
|
||||
progressHandler(
|
||||
Math.floor((progressEvent.loaded / progressEvent.total) * 100),
|
||||
);
|
||||
}
|
||||
});
|
||||
|
||||
// Use the project's baseApiUrl to construct the endpoint consistently
|
||||
const apiUrl = `${baseApiUrl('1.0')}chats/${conversationId}/attachments/backend-upload/`;
|
||||
xhr.open('POST', apiUrl);
|
||||
|
||||
// Add authentication headers following the project's pattern
|
||||
xhr.withCredentials = true;
|
||||
if (csrfToken) {
|
||||
xhr.setRequestHeader('X-CSRFToken', csrfToken);
|
||||
}
|
||||
|
||||
xhr.send(formData);
|
||||
});
|
||||
|
||||
export const useUploadFile = (conversationId: string) => {
|
||||
const {
|
||||
mutateAsync: createConversationAttachment,
|
||||
@@ -54,8 +124,25 @@ export const useUploadFile = (conversationId: string) => {
|
||||
error: errorAttachment,
|
||||
} = useCreateConversationAttachment();
|
||||
|
||||
const { data: conf } = useConfig();
|
||||
|
||||
const uploadFile = useCallback(
|
||||
async (file: File, progressHandler?: (progress: number) => void) => {
|
||||
// Backend mode backend_to_s3 file is sent to API backend
|
||||
if (conf?.FILE_UPLOAD_MODE === 'backend_to_s3') {
|
||||
// Upload file to backend (backend handles S3 storage, MIME detection, and malware scanning)
|
||||
const finalAttachment = await uploadFileToBackend(
|
||||
conversationId,
|
||||
file,
|
||||
(progress) => {
|
||||
progressHandler?.(progress);
|
||||
},
|
||||
);
|
||||
|
||||
return `/media-key/${finalAttachment.key}`;
|
||||
}
|
||||
|
||||
// Presigned URL mode (default): frontend uploads directly to S3
|
||||
const attachment = await createConversationAttachment({
|
||||
conversationId,
|
||||
content_type: file.type,
|
||||
@@ -83,7 +170,7 @@ export const useUploadFile = (conversationId: string) => {
|
||||
|
||||
return `/media-key/${attachment.key}`;
|
||||
},
|
||||
[createConversationAttachment, conversationId],
|
||||
[createConversationAttachment, conversationId, conf],
|
||||
);
|
||||
|
||||
return {
|
||||
|
||||
@@ -8,6 +8,7 @@ export const CONFIG = {
|
||||
'document-upload': 'enabled',
|
||||
'web-search': 'enabled',
|
||||
},
|
||||
FILE_UPLOAD_MODE: 'presigned_url',
|
||||
FRONTEND_CSS_URL: null,
|
||||
FRONTEND_HOMEPAGE_FEATURE_ENABLED: true,
|
||||
FRONTEND_THEME: null,
|
||||
|
||||
Reference in New Issue
Block a user