Merge pull request 'Cowork E: drop files into the chat (PDF, text, CSV, images)' (#66) from feat/cowork-file-drop into main
This commit was merged in pull request #66.
This commit is contained in:
@@ -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,87 @@ 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 ""
|
||||
# PyMuPDF is CPU-bound; keep it off the event loop so other streams don't stall.
|
||||
return name, await asyncio.to_thread(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 = await asyncio.to_thread(extract_text, file_bytes, ext)
|
||||
except FileExtractionError as e:
|
||||
return jsonify({"error": str(e)}), 400
|
||||
if not text.strip():
|
||||
message = (
|
||||
"This PDF has no text layer (scanned?), so Simba can't read it yet."
|
||||
if ext == "pdf"
|
||||
else "This file is empty, so there's nothing for Simba to read."
|
||||
)
|
||||
return jsonify({"error": message}), 422
|
||||
|
||||
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/<path:image_key>")
|
||||
@jwt_access_required
|
||||
async def serve_image(image_key: str):
|
||||
@@ -219,12 +311,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,
|
||||
|
||||
@@ -146,6 +146,8 @@ pub fn run() {
|
||||
.title("Simba")
|
||||
.inner_size(1200.0, 820.0)
|
||||
.min_inner_size(420.0, 520.0)
|
||||
// Let HTML5 drag-and-drop reach the webview (file drop in chat).
|
||||
.disable_drag_drop_handler()
|
||||
.build()?;
|
||||
build_tray(app.handle())?;
|
||||
register_hotkey(app.handle())?;
|
||||
|
||||
@@ -51,6 +51,12 @@ interface Conversation {
|
||||
user_id?: string;
|
||||
}
|
||||
|
||||
export type UploadResult =
|
||||
| { kind: "image"; image_key: string }
|
||||
| { kind: "file"; file_key: string; name: string; chars: number };
|
||||
|
||||
export class UploadError extends Error {}
|
||||
|
||||
interface QueryResponse {
|
||||
response: string;
|
||||
}
|
||||
@@ -182,16 +188,16 @@ class ConversationService {
|
||||
return await response.json();
|
||||
}
|
||||
|
||||
async uploadImage(
|
||||
async uploadAttachment(
|
||||
file: File,
|
||||
conversationId: string,
|
||||
): Promise<{ image_key: string }> {
|
||||
): Promise<UploadResult> {
|
||||
const formData = new FormData();
|
||||
formData.append("file", file);
|
||||
formData.append("conversation_id", conversationId);
|
||||
|
||||
const response = await userService.fetchWithAuth(
|
||||
`${this.conversationBaseUrl}/upload-image`,
|
||||
`${this.conversationBaseUrl}/upload`,
|
||||
{
|
||||
method: "POST",
|
||||
body: formData,
|
||||
@@ -200,8 +206,8 @@ class ConversationService {
|
||||
);
|
||||
|
||||
if (!response.ok) {
|
||||
const data = await response.json();
|
||||
throw new Error(data.error || "Failed to upload image");
|
||||
const data = await response.json().catch(() => ({}));
|
||||
throw new UploadError(data.error || "Failed to upload file");
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
@@ -224,11 +230,15 @@ class ConversationService {
|
||||
onEvent: SSEEventCallback,
|
||||
signal?: AbortSignal,
|
||||
imageKey?: string,
|
||||
fileKey?: string,
|
||||
): Promise<void> {
|
||||
const body: Record<string, string> = { query, conversation_id };
|
||||
if (imageKey) {
|
||||
body.image_key = imageKey;
|
||||
}
|
||||
if (fileKey) {
|
||||
body.file_key = fileKey;
|
||||
}
|
||||
|
||||
const response = await userService.fetchWithAuth(
|
||||
`${this.conversationBaseUrl}/stream-query`,
|
||||
|
||||
@@ -6,6 +6,7 @@ import { AnswerBubble } from "./AnswerBubble";
|
||||
import { ActivityTimeline } from "./ActivityTimeline";
|
||||
import { ThinkingBubble } from "./ThinkingBubble";
|
||||
import { MessageInput } from "./MessageInput";
|
||||
import { DropOverlay } from "./DropOverlay";
|
||||
import { SidebarContent } from "./SidebarContent";
|
||||
import { AdminPanel } from "./AdminPanel";
|
||||
import { ScheduledMessagesPanel } from "./ScheduledMessagesPanel";
|
||||
@@ -84,8 +85,11 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
|
||||
streaming,
|
||||
thinking,
|
||||
messagesLoading,
|
||||
pendingImage,
|
||||
setPendingImage,
|
||||
pendingAttachment,
|
||||
attachFile,
|
||||
clearAttachment,
|
||||
attachmentError,
|
||||
clearAttachmentError,
|
||||
sendMessage,
|
||||
stopGeneration,
|
||||
error,
|
||||
@@ -141,9 +145,6 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
|
||||
[deleteConversation],
|
||||
);
|
||||
|
||||
const handleImageSelect = useCallback((file: File) => setPendingImage(file), [setPendingImage]);
|
||||
const handleClearImage = useCallback(() => setPendingImage(null), [setPendingImage]);
|
||||
|
||||
const handleLogout = () => {
|
||||
localStorage.removeItem("access_token");
|
||||
localStorage.removeItem("refresh_token");
|
||||
@@ -223,7 +224,8 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
|
||||
{showAdminPanel && <AdminPanel onClose={() => setShowAdminPanel(false)} />}
|
||||
{showScheduler && <ScheduledMessagesPanel onClose={() => setShowScheduler(false)} />}
|
||||
|
||||
<div className="flex-1 flex flex-col h-full overflow-hidden min-w-0">
|
||||
<div className="relative flex-1 flex flex-col h-full overflow-hidden min-w-0">
|
||||
<DropOverlay onDrop={attachFile} />
|
||||
<header className="md:hidden flex items-center justify-between px-4 py-3 bg-warm-white border-b border-sand-light/60">
|
||||
<div className="flex items-center gap-2">
|
||||
<img src={catIcon} alt="Simba" className="w-12 h-12" />
|
||||
@@ -302,9 +304,11 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
|
||||
onToggleSimbaMode={handleToggleSimbaMode}
|
||||
isLoading={isLoading}
|
||||
onStop={stopGeneration}
|
||||
pendingImage={pendingImage}
|
||||
onImageSelect={handleImageSelect}
|
||||
onClearImage={handleClearImage}
|
||||
pendingAttachment={pendingAttachment}
|
||||
onAttach={attachFile}
|
||||
onClearAttachment={clearAttachment}
|
||||
attachmentError={attachmentError}
|
||||
onDismissAttachmentError={clearAttachmentError}
|
||||
autoFocus
|
||||
/>
|
||||
</div>
|
||||
@@ -431,9 +435,11 @@ export const ChatScreen = ({ setAuthenticated, isAdmin }: ChatScreenProps) => {
|
||||
onToggleSimbaMode={handleToggleSimbaMode}
|
||||
isLoading={isLoading}
|
||||
onStop={stopGeneration}
|
||||
pendingImage={pendingImage}
|
||||
onImageSelect={handleImageSelect}
|
||||
onClearImage={handleClearImage}
|
||||
pendingAttachment={pendingAttachment}
|
||||
onAttach={attachFile}
|
||||
onClearAttachment={clearAttachment}
|
||||
attachmentError={attachmentError}
|
||||
onDismissAttachmentError={clearAttachmentError}
|
||||
/>
|
||||
</div>
|
||||
</footer>
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
import { useEffect, useRef, useState } from "react";
|
||||
import { Paperclip } from "lucide-react";
|
||||
|
||||
const hasFiles = (e: DragEvent) =>
|
||||
Array.from(e.dataTransfer?.types ?? []).includes("Files");
|
||||
|
||||
type DropOverlayProps = {
|
||||
onDrop: (file: File) => void;
|
||||
};
|
||||
|
||||
/**
|
||||
* Listens for file drags anywhere in the window and covers its positioned
|
||||
* parent (the chat column, beside the artifact panel) with a "Drop to attach"
|
||||
* target. Below lg the panel is a full-screen sheet, so it covers the window
|
||||
* and sits above the sheet. Drags without files (text, links)
|
||||
* are ignored so they keep their normal behaviour.
|
||||
*/
|
||||
export const DropOverlay = ({ onDrop }: DropOverlayProps) => {
|
||||
const [active, setActive] = useState(false);
|
||||
// dragenter/dragleave fire for every element crossed; count to know when
|
||||
// the drag has really left the window.
|
||||
const depthRef = useRef(0);
|
||||
const onDropRef = useRef(onDrop);
|
||||
onDropRef.current = onDrop;
|
||||
|
||||
useEffect(() => {
|
||||
const handleEnter = (e: DragEvent) => {
|
||||
if (!hasFiles(e)) return;
|
||||
e.preventDefault();
|
||||
depthRef.current += 1;
|
||||
setActive(true);
|
||||
};
|
||||
const handleOver = (e: DragEvent) => {
|
||||
if (!hasFiles(e)) return;
|
||||
// Required for the drop event to fire (and to stop the browser
|
||||
// navigating to the file).
|
||||
e.preventDefault();
|
||||
if (e.dataTransfer) e.dataTransfer.dropEffect = "copy";
|
||||
};
|
||||
const handleLeave = (e: DragEvent) => {
|
||||
if (!hasFiles(e)) return;
|
||||
depthRef.current = Math.max(0, depthRef.current - 1);
|
||||
if (depthRef.current === 0) setActive(false);
|
||||
};
|
||||
const handleDrop = (e: DragEvent) => {
|
||||
if (!hasFiles(e)) return;
|
||||
e.preventDefault();
|
||||
depthRef.current = 0;
|
||||
setActive(false);
|
||||
const file = e.dataTransfer?.files?.[0];
|
||||
if (file) onDropRef.current(file);
|
||||
};
|
||||
|
||||
window.addEventListener("dragenter", handleEnter);
|
||||
window.addEventListener("dragover", handleOver);
|
||||
window.addEventListener("dragleave", handleLeave);
|
||||
window.addEventListener("drop", handleDrop);
|
||||
return () => {
|
||||
window.removeEventListener("dragenter", handleEnter);
|
||||
window.removeEventListener("dragover", handleOver);
|
||||
window.removeEventListener("dragleave", handleLeave);
|
||||
window.removeEventListener("drop", handleDrop);
|
||||
};
|
||||
}, []);
|
||||
|
||||
if (!active) return null;
|
||||
|
||||
return (
|
||||
<div
|
||||
className="pointer-events-none fixed lg:absolute inset-0 z-[45] flex items-center justify-center bg-cream/80 backdrop-blur-sm p-3 md:p-5"
|
||||
data-testid="drop-overlay"
|
||||
>
|
||||
<div className="w-full h-full rounded-3xl border-2 border-dashed border-amber-glow/50 bg-amber-pale/70 flex flex-col items-center justify-center gap-3 message-enter">
|
||||
<div className="w-14 h-14 rounded-2xl bg-warm-white border border-sand shadow-md shadow-sand/30 text-amber-glow flex items-center justify-center">
|
||||
<Paperclip size={24} />
|
||||
</div>
|
||||
<p
|
||||
className="text-lg font-bold text-charcoal"
|
||||
style={{ fontFamily: "var(--font-display)" }}
|
||||
>
|
||||
Drop to attach
|
||||
</p>
|
||||
<p className="text-xs text-warm-gray">Images, PDF, TXT, MD or CSV · up to 10 MB</p>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
@@ -1,6 +1,12 @@
|
||||
import React, { useEffect, useMemo, useRef } from "react";
|
||||
import { ArrowUp, ImagePlus, Square, X } from "lucide-react";
|
||||
import { ArrowUp, FileSpreadsheet, FileText, Paperclip, Square, X } from "lucide-react";
|
||||
import { cn } from "../lib/utils";
|
||||
import {
|
||||
ATTACHMENT_ACCEPT,
|
||||
attachmentKind,
|
||||
fileExtension,
|
||||
formatBytes,
|
||||
} from "../lib/attachments";
|
||||
import { Textarea } from "./ui/textarea";
|
||||
|
||||
type MessageInputProps = {
|
||||
@@ -12,9 +18,11 @@ type MessageInputProps = {
|
||||
query: string;
|
||||
isLoading: boolean;
|
||||
onStop: () => void;
|
||||
pendingImage: File | null;
|
||||
onImageSelect: (file: File) => void;
|
||||
onClearImage: () => void;
|
||||
pendingAttachment: File | null;
|
||||
onAttach: (file: File) => void;
|
||||
onClearAttachment: () => void;
|
||||
attachmentError: string | null;
|
||||
onDismissAttachmentError: () => void;
|
||||
autoFocus?: boolean;
|
||||
};
|
||||
|
||||
@@ -28,15 +36,23 @@ export const MessageInput = React.memo(
|
||||
onToggleSimbaMode,
|
||||
isLoading,
|
||||
onStop,
|
||||
pendingImage,
|
||||
onImageSelect,
|
||||
onClearImage,
|
||||
pendingAttachment,
|
||||
onAttach,
|
||||
onClearAttachment,
|
||||
attachmentError,
|
||||
onDismissAttachmentError,
|
||||
autoFocus,
|
||||
}: MessageInputProps) => {
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
const textareaRef = useRef<HTMLTextAreaElement>(null);
|
||||
|
||||
// Create blob URL once per file, revoke on cleanup
|
||||
const pendingImage =
|
||||
pendingAttachment && attachmentKind(pendingAttachment) === "image"
|
||||
? pendingAttachment
|
||||
: null;
|
||||
const pendingFile = pendingAttachment && !pendingImage ? pendingAttachment : null;
|
||||
|
||||
// Create blob URL once per image, revoke on cleanup
|
||||
const previewUrl = useMemo(
|
||||
() => (pendingImage ? URL.createObjectURL(pendingImage) : null),
|
||||
[pendingImage],
|
||||
@@ -59,7 +75,7 @@ export const MessageInput = React.memo(
|
||||
const handleFileChange = (e: React.ChangeEvent<HTMLInputElement>) => {
|
||||
const file = e.target.files?.[0];
|
||||
if (file) {
|
||||
onImageSelect(file);
|
||||
onAttach(file);
|
||||
}
|
||||
// Reset so the same file can be re-selected
|
||||
e.target.value = "";
|
||||
@@ -69,19 +85,21 @@ export const MessageInput = React.memo(
|
||||
if (isLoading) return;
|
||||
const items = e.clipboardData?.items;
|
||||
if (!items) return;
|
||||
for (const item of items) {
|
||||
for (const item of Array.from(items)) {
|
||||
if (item.type.startsWith("image/")) {
|
||||
const file = item.getAsFile();
|
||||
if (file) {
|
||||
e.preventDefault();
|
||||
onImageSelect(file);
|
||||
onAttach(file);
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const canSend = !isLoading && Boolean(query.trim() || pendingImage);
|
||||
const canSend = !isLoading && Boolean(query.trim() || pendingAttachment);
|
||||
const FileIcon =
|
||||
pendingFile && fileExtension(pendingFile.name) === "csv" ? FileSpreadsheet : FileText;
|
||||
|
||||
return (
|
||||
<div
|
||||
@@ -102,7 +120,8 @@ export const MessageInput = React.memo(
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClearImage}
|
||||
onClick={onClearAttachment}
|
||||
aria-label="Remove attachment"
|
||||
className="absolute -top-1.5 -right-1.5 w-5 h-5 rounded-full bg-charcoal text-white flex items-center justify-center hover:bg-charcoal/80 transition-colors cursor-pointer"
|
||||
>
|
||||
<X size={12} />
|
||||
@@ -111,6 +130,56 @@ export const MessageInput = React.memo(
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* File chip */}
|
||||
{pendingFile && (
|
||||
<div className="px-3 pt-3">
|
||||
<div
|
||||
className="relative inline-flex items-center gap-2.5 max-w-full rounded-xl border border-sand bg-cream pl-2 pr-7 py-2 message-enter"
|
||||
data-testid="attachment-chip"
|
||||
>
|
||||
<div className="w-8 h-8 shrink-0 rounded-lg bg-amber-pale text-amber-dark flex items-center justify-center">
|
||||
<FileIcon size={16} />
|
||||
</div>
|
||||
<div className="min-w-0">
|
||||
<div className="text-xs font-medium text-charcoal truncate max-w-[14rem] sm:max-w-xs">
|
||||
{pendingFile.name}
|
||||
</div>
|
||||
<div className="text-[11px] text-warm-gray">
|
||||
{fileExtension(pendingFile.name).toUpperCase() || "File"} ·{" "}
|
||||
{formatBytes(pendingFile.size)}
|
||||
</div>
|
||||
</div>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClearAttachment}
|
||||
aria-label="Remove attachment"
|
||||
className="absolute top-1.5 right-1.5 w-5 h-5 rounded-full text-warm-gray flex items-center justify-center hover:text-charcoal hover:bg-cream-dark transition-colors cursor-pointer"
|
||||
>
|
||||
<X size={12} />
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{attachmentError && (
|
||||
<div className="px-3 pt-3">
|
||||
<div
|
||||
role="alert"
|
||||
className="inline-flex items-center gap-2 px-3 py-1.5 rounded-full bg-red-50 border border-red-200 text-xs text-red-600 message-enter"
|
||||
>
|
||||
<span>{attachmentError}</span>
|
||||
<button
|
||||
type="button"
|
||||
onClick={onDismissAttachmentError}
|
||||
className="text-red-400 hover:text-red-600 transition-colors cursor-pointer"
|
||||
aria-label="Dismiss error"
|
||||
>
|
||||
<X size={13} />
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Textarea */}
|
||||
<Textarea
|
||||
ref={textareaRef}
|
||||
@@ -128,7 +197,7 @@ export const MessageInput = React.memo(
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
type="file"
|
||||
accept="image/*"
|
||||
accept={ATTACHMENT_ACCEPT}
|
||||
onChange={handleFileChange}
|
||||
className="hidden"
|
||||
/>
|
||||
@@ -150,11 +219,13 @@ export const MessageInput = React.memo(
|
||||
</span>
|
||||
</button>
|
||||
|
||||
{/* Image attach button */}
|
||||
{/* Attach button */}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => fileInputRef.current?.click()}
|
||||
disabled={isLoading}
|
||||
aria-label="Attach a file"
|
||||
title="Attach an image, PDF, TXT, MD or CSV"
|
||||
className={cn(
|
||||
"w-7 h-7 rounded-lg flex items-center justify-center transition-all cursor-pointer",
|
||||
isLoading
|
||||
@@ -162,7 +233,7 @@ export const MessageInput = React.memo(
|
||||
: "text-warm-gray hover:text-charcoal hover:bg-cream-dark",
|
||||
)}
|
||||
>
|
||||
<ImagePlus size={16} />
|
||||
<Paperclip size={16} />
|
||||
</button>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import React from "react";
|
||||
import { FileSpreadsheet, FileText } from "lucide-react";
|
||||
import { cn } from "../lib/utils";
|
||||
import { fileExtension, splitAttachedFile } from "../lib/attachments";
|
||||
import { usePresignedUrl } from "../hooks/usePresignedUrl";
|
||||
|
||||
type QuestionBubbleProps = {
|
||||
@@ -9,6 +11,8 @@ type QuestionBubbleProps = {
|
||||
|
||||
export const QuestionBubble = React.memo(({ text, image_key }: QuestionBubbleProps) => {
|
||||
const { imageUrl, imageError } = usePresignedUrl(image_key);
|
||||
const { body, fileName } = splitAttachedFile(text);
|
||||
const FileIcon = fileName && fileExtension(fileName) === "csv" ? FileSpreadsheet : FileText;
|
||||
|
||||
return (
|
||||
<div className="flex justify-end message-enter">
|
||||
@@ -32,7 +36,21 @@ export const QuestionBubble = React.memo(({ text, image_key }: QuestionBubblePro
|
||||
className="max-w-full rounded-xl mb-2"
|
||||
/>
|
||||
)}
|
||||
{text}
|
||||
{fileName && (
|
||||
<div
|
||||
className={cn(
|
||||
"flex items-center gap-1.5 w-fit max-w-full rounded-lg",
|
||||
"bg-warm-white/70 border border-leaf-light/70 px-2 py-1",
|
||||
"text-xs text-leaf-dark",
|
||||
body && "mb-2",
|
||||
)}
|
||||
data-testid="message-attachment"
|
||||
>
|
||||
<FileIcon size={13} className="shrink-0" />
|
||||
<span className="truncate">{fileName}</span>
|
||||
</div>
|
||||
)}
|
||||
{body}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
|
||||
@@ -2,9 +2,11 @@ import { useState, useCallback, useEffect, useRef } from "react";
|
||||
import { useMatch } from "react-router-dom";
|
||||
import {
|
||||
conversationService,
|
||||
UploadError,
|
||||
type ToolInput,
|
||||
} from "../api/conversationService";
|
||||
import type { ArtifactRef } from "../api/artifactTypes";
|
||||
import { validateAttachment, withAttachedFile } from "../lib/attachments";
|
||||
import type { Conversation } from "./useConversations";
|
||||
|
||||
export type ToolStatus = "running" | "ok" | "error";
|
||||
@@ -63,7 +65,9 @@ export function useChat({
|
||||
const [messages, setMessages] = useState<Message[]>([]);
|
||||
const [isLoading, setIsLoading] = useState(false);
|
||||
const [messagesLoading, setMessagesLoading] = useState(false);
|
||||
const [pendingImage, setPendingImage] = useState<File | null>(null);
|
||||
// One image or text-like file to send with the next message.
|
||||
const [pendingAttachment, setPendingAttachment] = useState<File | null>(null);
|
||||
const [attachmentError, setAttachmentError] = useState<string | null>(null);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [stalled, setStalled] = useState(false);
|
||||
const [streaming, setStreaming] = useState(false);
|
||||
@@ -164,7 +168,7 @@ export function useChat({
|
||||
|
||||
const sendMessage = useCallback(
|
||||
async (query: string, simbaMode: boolean) => {
|
||||
if ((!query.trim() && !pendingImage) || isLoading) return;
|
||||
if ((!query.trim() && !pendingAttachment) || isLoading) return;
|
||||
|
||||
let activeId = conversationId;
|
||||
if (!activeId) {
|
||||
@@ -174,7 +178,7 @@ export function useChat({
|
||||
activeId = created.id;
|
||||
}
|
||||
|
||||
const imageFile = pendingImage;
|
||||
const attachment = pendingAttachment;
|
||||
|
||||
setError(null);
|
||||
setStalled(false);
|
||||
@@ -185,7 +189,8 @@ export function useChat({
|
||||
streamingActiveRef.current = false;
|
||||
setStreaming(false);
|
||||
updateMessages((prev) => prev.concat([{ text: query, speaker: "user" }]));
|
||||
setPendingImage(null);
|
||||
setPendingAttachment(null);
|
||||
setAttachmentError(null);
|
||||
setIsLoading(true);
|
||||
|
||||
if (simbaMode) {
|
||||
@@ -203,19 +208,27 @@ export function useChat({
|
||||
|
||||
try {
|
||||
let imageKey: string | undefined;
|
||||
if (imageFile) {
|
||||
const uploadResult = await conversationService.uploadImage(
|
||||
imageFile,
|
||||
let fileKey: string | undefined;
|
||||
if (attachment) {
|
||||
const uploadResult = await conversationService.uploadAttachment(
|
||||
attachment,
|
||||
activeId,
|
||||
);
|
||||
imageKey = uploadResult.image_key;
|
||||
let patch: Partial<TextMessage>;
|
||||
if (uploadResult.kind === "image") {
|
||||
imageKey = uploadResult.image_key;
|
||||
patch = { image_key: imageKey };
|
||||
} else {
|
||||
fileKey = uploadResult.file_key;
|
||||
patch = { text: withAttachedFile(query, uploadResult.name) };
|
||||
}
|
||||
|
||||
updateMessages((prev) => {
|
||||
const updated = [...prev];
|
||||
for (let i = updated.length - 1; i >= 0; i--) {
|
||||
const msg = updated[i];
|
||||
if (msg.speaker === "user") {
|
||||
updated[i] = { ...msg, image_key: imageKey };
|
||||
updated[i] = { ...msg, ...patch };
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -342,10 +355,13 @@ export function useChat({
|
||||
},
|
||||
abortController.signal,
|
||||
imageKey,
|
||||
fileKey,
|
||||
);
|
||||
} catch (err) {
|
||||
if (err instanceof Error && err.name === "AbortError") {
|
||||
// User cancelled — not an error to surface.
|
||||
} else if (err instanceof UploadError) {
|
||||
if (isMountedRef.current) setError(`Couldn't attach that file: ${err.message}`);
|
||||
} else {
|
||||
console.error("Failed to send query:", err);
|
||||
if (
|
||||
@@ -383,7 +399,7 @@ export function useChat({
|
||||
}
|
||||
},
|
||||
[
|
||||
pendingImage,
|
||||
pendingAttachment,
|
||||
isLoading,
|
||||
conversationId,
|
||||
createConversation,
|
||||
@@ -400,6 +416,19 @@ export function useChat({
|
||||
|
||||
const clearError = useCallback(() => setError(null), []);
|
||||
|
||||
const attachFile = useCallback((file: File) => {
|
||||
const problem = validateAttachment(file);
|
||||
setAttachmentError(problem);
|
||||
if (!problem) setPendingAttachment(file);
|
||||
}, []);
|
||||
|
||||
const clearAttachment = useCallback(() => {
|
||||
setPendingAttachment(null);
|
||||
setAttachmentError(null);
|
||||
}, []);
|
||||
|
||||
const clearAttachmentError = useCallback(() => setAttachmentError(null), []);
|
||||
|
||||
const retryStalled = useCallback(() => {
|
||||
const q = stalledQueryRef.current;
|
||||
if (!q) return;
|
||||
@@ -415,8 +444,11 @@ export function useChat({
|
||||
streaming,
|
||||
thinking,
|
||||
messagesLoading,
|
||||
pendingImage,
|
||||
setPendingImage,
|
||||
pendingAttachment,
|
||||
attachFile,
|
||||
clearAttachment,
|
||||
attachmentError,
|
||||
clearAttachmentError,
|
||||
sendMessage,
|
||||
stopGeneration,
|
||||
error,
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
// What can be attached to a message: images go to the vision model, text-like
|
||||
// files are extracted server-side and appended to the turn.
|
||||
|
||||
export const ATTACHMENT_ACCEPT = "image/*,.pdf,.txt,.md,.markdown,.csv";
|
||||
export const MAX_ATTACHMENT_BYTES = 10 * 1024 * 1024;
|
||||
|
||||
const IMAGE_TYPES = new Set([
|
||||
"image/jpeg",
|
||||
"image/png",
|
||||
"image/webp",
|
||||
"image/heic",
|
||||
"image/heif",
|
||||
]);
|
||||
const FILE_EXTENSIONS = new Set(["pdf", "txt", "md", "markdown", "csv"]);
|
||||
const FILE_TYPES = new Set([
|
||||
"application/pdf",
|
||||
"text/plain",
|
||||
"text/markdown",
|
||||
"text/x-markdown",
|
||||
"text/csv",
|
||||
]);
|
||||
|
||||
export type AttachmentKind = "image" | "file";
|
||||
|
||||
export const fileExtension = (name: string) =>
|
||||
name.includes(".") ? name.split(".").pop()!.toLowerCase() : "";
|
||||
|
||||
export function attachmentKind(file: File): AttachmentKind | null {
|
||||
if (IMAGE_TYPES.has(file.type)) return "image";
|
||||
if (FILE_TYPES.has(file.type) || FILE_EXTENSIONS.has(fileExtension(file.name))) {
|
||||
return "file";
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Returns an error message if the file can't be attached, else null. */
|
||||
export function validateAttachment(file: File): string | null {
|
||||
if (!attachmentKind(file)) {
|
||||
return `Can't attach “${file.name}”. Use an image, PDF, TXT, MD or CSV file.`;
|
||||
}
|
||||
if (file.size > MAX_ATTACHMENT_BYTES) {
|
||||
return `“${file.name}” is larger than 10 MB.`;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
export function formatBytes(bytes: number): string {
|
||||
if (bytes < 1024) return `${bytes} B`;
|
||||
if (bytes < 1024 * 1024) return `${Math.round(bytes / 1024)} KB`;
|
||||
return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
|
||||
}
|
||||
|
||||
// Saved user messages record an attached file as a trailing "📎 name" line.
|
||||
const ATTACHED_FILE_RE = /(?:^|\n\n)📎 ([^\n]+)$/u;
|
||||
|
||||
export const withAttachedFile = (text: string, name: string) =>
|
||||
text ? `${text}\n\n📎 ${name}` : `📎 ${name}`;
|
||||
|
||||
export function splitAttachedFile(text: string): { body: string; fileName: string | null } {
|
||||
const match = text.match(ATTACHED_FILE_RE);
|
||||
if (!match) return { body: text, fileName: null };
|
||||
return { body: text.slice(0, match.index), fileName: match[1] };
|
||||
}
|
||||
@@ -20,8 +20,16 @@
|
||||
// (all v1 kinds) | "broken" (GET /api/artifacts/<id> 500s) | "mystery"
|
||||
// (unknown kind → JSON fallback)
|
||||
// anything else: two sequential calls, then a streamed answer
|
||||
// Conversations "seeded" (timeline) and "everything" (one saved tool message
|
||||
// per artifact kind) exist up front.
|
||||
// any query sent with a file_key: reads the attached file (by name)
|
||||
// Conversations "seeded" (timeline), "everything" (one saved tool message
|
||||
// per artifact kind) and "attachment" (a saved 📎 file message) exist up front.
|
||||
//
|
||||
// Uploads (POST /api/conversation/upload) follow the contract: JPEG/PNG/WebP/
|
||||
// HEIC -> {kind:"image"}, PDF/TXT/MD/CSV (by MIME or extension) ->
|
||||
// {kind:"file"}, anything else 415, files named *scan*.pdf 422 (no text
|
||||
// layer). Uploaded images are served back via /api/conversation/image/<key>.
|
||||
// GET /__mock/requests returns recorded API request bodies (multipart files
|
||||
// summarised as {filename, type, size}), newest last.
|
||||
// artifacts/<kind>.json holds one ToolArtifact per v1 kind, matching what the
|
||||
// backend tools emit. Any other *.json in this directory with `kind` and
|
||||
// `data` is loaded as an artifact of "seeded" (`id` defaults to the file name).
|
||||
@@ -62,6 +70,10 @@ const uuid = () => crypto.randomUUID();
|
||||
const artifacts = new Map();
|
||||
/** @type {Map<string, {id, name, created_at, updated_at, messages: any[]}>} */
|
||||
const conversations = new Map();
|
||||
/** Uploaded objects by S3-style key: {bytes, type, name} */
|
||||
const blobs = new Map();
|
||||
/** Recorded API requests: {method, path, body} */
|
||||
const requests = [];
|
||||
|
||||
function addArtifact(conversationId, tool, a) {
|
||||
const record = {
|
||||
@@ -278,8 +290,15 @@ const LABELS = {
|
||||
};
|
||||
const label = (tool) => LABELS[tool] ?? `Using ${tool}…`;
|
||||
|
||||
function seedAttachment() {
|
||||
const c = newConversation("attachment", "Vet bill questions");
|
||||
pushMessage(c, { speaker: "user", text: "How much was the dental cleaning?\n\n📎 vet-invoice-2026-09.pdf" });
|
||||
pushMessage(c, { speaker: "simba", text: "The dental cleaning on the September invoice was **$412.50**." });
|
||||
}
|
||||
|
||||
seed();
|
||||
seedEverything();
|
||||
seedAttachment();
|
||||
loadFixtureArtifacts();
|
||||
|
||||
/** Artifact-panel scenarios: one tool call per artifact, by keyword. */
|
||||
@@ -304,11 +323,20 @@ function artifactScenario(q) {
|
||||
* Each step is [delayMs, event | fn(ctx) => event]. Tool steps are expanded
|
||||
* into contract-shaped events; ctx persists tool messages and artifacts.
|
||||
*/
|
||||
function scenario(query) {
|
||||
function scenario(query, file) {
|
||||
const q = query.toLowerCase();
|
||||
const answer = (text) =>
|
||||
text.match(/\S+\s*/g).map((w) => [60, { type: "content", delta: w }]);
|
||||
|
||||
if (file) {
|
||||
return [
|
||||
[200, { type: "thinking_delta", delta: `Reading the attached ${file.name}…` }],
|
||||
[300, { start: "f1", tool: "simba_search", input: { query: (query || file.name).slice(0, 200) } }],
|
||||
[600, { end: "f1" }],
|
||||
...answer(`I read **${file.name}**. It looks like a vet invoice; the dental cleaning came to $412.50.`),
|
||||
];
|
||||
}
|
||||
|
||||
if (q.includes("plain")) {
|
||||
return [[300, { type: "thinking" }], ...answer("Meow. Nothing to look up for that one.")];
|
||||
}
|
||||
@@ -363,8 +391,17 @@ function scenario(query) {
|
||||
|
||||
async function streamQuery(req, res, body) {
|
||||
const c = conversations.get(body.conversation_id) ?? newConversation(body.conversation_id);
|
||||
pushMessage(c, { speaker: "user", text: body.query, image_key: body.image_key ?? null });
|
||||
if (c.name === "New Conversation") c.name = body.query.slice(0, 40);
|
||||
let saved = body.query ?? "";
|
||||
let file = null;
|
||||
if (body.file_key) {
|
||||
file = blobs.get(body.file_key);
|
||||
if (!body.file_key.startsWith(`conversations/${c.id}/files/`) || !file) {
|
||||
return json(res, 400, { error: "invalid file_key" });
|
||||
}
|
||||
saved = `${saved}\n\n📎 ${file.name}`.trim();
|
||||
}
|
||||
pushMessage(c, { speaker: "user", text: saved, image_key: body.image_key ?? null });
|
||||
if (c.name === "New Conversation") c.name = (body.query || file?.name || "Attachment").slice(0, 40);
|
||||
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "text/event-stream",
|
||||
@@ -377,7 +414,7 @@ async function streamQuery(req, res, body) {
|
||||
|
||||
const runs = new Map(); // key -> {run_id, tool, started}
|
||||
let text = "";
|
||||
for (const [delay, step] of scenario(body.query)) {
|
||||
for (const [delay, step] of scenario(body.query ?? "", file)) {
|
||||
await sleep(delay);
|
||||
if (closed) return;
|
||||
if (step.start) {
|
||||
@@ -428,24 +465,100 @@ const json = (res, status, body) => {
|
||||
res.end(JSON.stringify(body));
|
||||
};
|
||||
|
||||
const readBody = (req) =>
|
||||
new Promise((resolve) => {
|
||||
let raw = "";
|
||||
req.on("data", (chunk) => (raw += chunk));
|
||||
req.on("end", () => {
|
||||
try {
|
||||
resolve(raw ? JSON.parse(raw) : {});
|
||||
} catch {
|
||||
resolve({});
|
||||
}
|
||||
});
|
||||
});
|
||||
async function readRaw(req) {
|
||||
const chunks = [];
|
||||
for await (const chunk of req) chunks.push(chunk);
|
||||
return Buffer.concat(chunks);
|
||||
}
|
||||
|
||||
/** Minimal multipart/form-data parser: {fields, files: {name: {filename, type, data}}}. */
|
||||
function parseMultipart(buf, contentType) {
|
||||
const boundary = /boundary=(?:"([^"]+)"|([^;]+))/.exec(contentType ?? "");
|
||||
const out = { fields: {}, files: {} };
|
||||
if (!boundary) return out;
|
||||
const delim = Buffer.from(`--${boundary[1] ?? boundary[2]}`);
|
||||
let pos = buf.indexOf(delim);
|
||||
while (pos !== -1) {
|
||||
const start = pos + delim.length + 2; // skip CRLF
|
||||
const next = buf.indexOf(delim, start);
|
||||
if (next === -1) break;
|
||||
const part = buf.subarray(start, next - 2); // strip trailing CRLF
|
||||
const headerEnd = part.indexOf("\r\n\r\n");
|
||||
const headers = part.subarray(0, headerEnd).toString("utf8");
|
||||
const data = part.subarray(headerEnd + 4);
|
||||
const name = /name="([^"]*)"/.exec(headers)?.[1];
|
||||
const filename = /filename="([^"]*)"/.exec(headers)?.[1];
|
||||
if (name && filename !== undefined) {
|
||||
const type = /Content-Type:\s*([^\r\n]+)/i.exec(headers)?.[1]?.trim() ?? "";
|
||||
out.files[name] = { filename, type, data };
|
||||
} else if (name) {
|
||||
out.fields[name] = data.toString("utf8");
|
||||
}
|
||||
pos = next;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/** Read the body once: JSON, or multipart form as {fields, files}. */
|
||||
async function readBody(req) {
|
||||
const raw = await readRaw(req);
|
||||
const type = req.headers["content-type"] ?? "";
|
||||
if (type.startsWith("multipart/form-data")) return { form: parseMultipart(raw, type) };
|
||||
try {
|
||||
return raw.length ? JSON.parse(raw.toString("utf8")) : {};
|
||||
} catch {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
const IMAGE_TYPES = new Set(["image/jpeg", "image/png", "image/webp", "image/heic", "image/heif"]);
|
||||
const FILE_TYPES = { "application/pdf": "pdf", "text/plain": "txt", "text/markdown": "md", "text/csv": "csv" };
|
||||
const FILE_EXTS = new Set(["pdf", "txt", "md", "csv"]);
|
||||
|
||||
function upload(res, form) {
|
||||
const file = form?.files.file;
|
||||
const conversationId = form?.fields.conversation_id;
|
||||
if (!file || !conversationId) return json(res, 400, { error: "file and conversation_id are required" });
|
||||
if (!conversations.has(conversationId)) return json(res, 404, { error: "conversation not found" });
|
||||
if (file.data.length > 10 * 1024 * 1024) return json(res, 413, { error: "File too large" });
|
||||
const type = file.type.split(";")[0].trim().toLowerCase();
|
||||
if (IMAGE_TYPES.has(type)) {
|
||||
const key = `conversations/${conversationId}/${uuid()}.${type.split("/")[1].replace("jpeg", "jpg")}`;
|
||||
blobs.set(key, { bytes: file.data, type, name: file.filename });
|
||||
return json(res, 200, { kind: "image", image_key: key });
|
||||
}
|
||||
const fromName = path.extname(file.filename).slice(1).toLowerCase();
|
||||
const ext = FILE_TYPES[type] ?? (FILE_EXTS.has(fromName) ? fromName : null);
|
||||
if (!ext) {
|
||||
return json(res, 415, { error: `Unsupported file type: ${type || file.filename}. Allowed: images, PDF, TXT, MD, CSV` });
|
||||
}
|
||||
if (ext === "pdf" && /scan/i.test(file.filename)) {
|
||||
return json(res, 422, { error: "This PDF has no text layer (scanned?), so Simba can't read it yet." });
|
||||
}
|
||||
const key = `conversations/${conversationId}/files/${uuid()}.${ext}`;
|
||||
blobs.set(key, { bytes: file.data, type, name: file.filename });
|
||||
return json(res, 200, { kind: "file", file_key: key, name: file.filename, chars: ext === "pdf" ? 1834 : file.data.length });
|
||||
}
|
||||
|
||||
const summary = (c) => ({ id: c.id, name: c.name, created_at: c.created_at, updated_at: c.updated_at });
|
||||
|
||||
async function api(req, res, url) {
|
||||
const p = url.pathname.replace(/\/+$/, "") || "/";
|
||||
const m = req.method;
|
||||
const body = m === "GET" ? {} : await readBody(req);
|
||||
const { form } = body;
|
||||
requests.push({
|
||||
method: m,
|
||||
path: url.pathname + url.search,
|
||||
body: form
|
||||
? {
|
||||
...form.fields,
|
||||
files: Object.fromEntries(
|
||||
Object.entries(form.files).map(([k, f]) => [k, { filename: f.filename, type: f.type, size: f.data.length }]),
|
||||
),
|
||||
}
|
||||
: body,
|
||||
});
|
||||
|
||||
if (p === "/api/user/me") {
|
||||
return json(res, 200, { id: "u1", username: "ryan", email: "ryan@example.com", is_admin: true });
|
||||
@@ -456,16 +569,18 @@ async function api(req, res, url) {
|
||||
if (p.startsWith("/api/scheduled")) return json(res, 200, []);
|
||||
|
||||
if (p === "/api/conversation/stream-query" && m === "POST") {
|
||||
return streamQuery(req, res, await readBody(req));
|
||||
return streamQuery(req, res, body);
|
||||
}
|
||||
if (p === "/api/conversation/query" && m === "POST") {
|
||||
return json(res, 200, { response: "meow (mock)" });
|
||||
}
|
||||
if (p === "/api/conversation/upload-image" || p === "/api/conversation/upload") {
|
||||
return json(res, 200, { kind: "image", image_key: `mock/${uuid()}.png` });
|
||||
if (p === "/api/conversation/upload" && m === "POST") return upload(res, form);
|
||||
if (p === "/api/conversation/upload-image") {
|
||||
return json(res, 200, { image_key: `mock/${uuid()}.png` });
|
||||
}
|
||||
if (p.startsWith("/api/conversation/image/")) {
|
||||
return json(res, 200, { url: "/favicon.ico" });
|
||||
const key = decodeURIComponent(p.slice("/api/conversation/image/".length));
|
||||
return json(res, 200, { url: blobs.has(key) ? `/__mock/blob/${encodeURIComponent(key)}` : "/favicon.ico" });
|
||||
}
|
||||
if (p === "/api/conversation") {
|
||||
if (m === "POST") return json(res, 200, summary(newConversation()));
|
||||
@@ -490,7 +605,7 @@ async function api(req, res, url) {
|
||||
return json(res, 200, { success: true });
|
||||
}
|
||||
if (m === "PATCH") {
|
||||
c.name = (await readBody(req)).name ?? c.name;
|
||||
c.name = body.name ?? c.name;
|
||||
return json(res, 200, summary(c));
|
||||
}
|
||||
return json(res, 200, {
|
||||
@@ -547,6 +662,13 @@ http
|
||||
.createServer(async (req, res) => {
|
||||
const url = new URL(req.url, `http://${req.headers.host}`);
|
||||
try {
|
||||
if (url.pathname === "/__mock/requests") return json(res, 200, requests);
|
||||
if (url.pathname.startsWith("/__mock/blob/")) {
|
||||
const blob = blobs.get(decodeURIComponent(url.pathname.slice("/__mock/blob/".length)));
|
||||
if (!blob) return json(res, 404, { error: "no blob" });
|
||||
res.writeHead(200, { "Content-Type": blob.type });
|
||||
return res.end(blob.bytes);
|
||||
}
|
||||
if (url.pathname.startsWith("/api/")) return await api(req, res, url);
|
||||
return serveStatic(res, decodeURIComponent(url.pathname));
|
||||
} catch (err) {
|
||||
|
||||
@@ -0,0 +1,282 @@
|
||||
"""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('<attached file "a\'b.txt">\n')
|
||||
assert block.endswith(f"{TRUNCATION_MARKER}\n</attached>")
|
||||
|
||||
|
||||
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_pdf_without_text_layer_is_422(self, app, headers, s3):
|
||||
doc = fitz.open()
|
||||
doc.new_page() # blank page: no text, like a scan without OCR
|
||||
resp = await _upload(app, headers, doc.tobytes(), "scan.pdf", "application/pdf")
|
||||
assert resp.status_code == 422
|
||||
assert "no text layer" in (await resp.get_json())["error"]
|
||||
s3.upload_file.assert_not_awaited()
|
||||
|
||||
async def test_empty_text_file_is_422(self, app, headers, s3):
|
||||
resp = await _upload(app, headers, b" \n", "empty.txt", "text/plain")
|
||||
assert resp.status_code == 422
|
||||
s3.upload_file.assert_not_awaited()
|
||||
|
||||
async def test_extraction_runs_off_the_event_loop(self, app, headers, s3):
|
||||
calls = []
|
||||
|
||||
async def fake_to_thread(fn, *args):
|
||||
calls.append(fn)
|
||||
return fn(*args)
|
||||
|
||||
with patch.object(conv.asyncio, "to_thread", fake_to_thread):
|
||||
resp = await _upload(
|
||||
app, headers, _pdf_bytes("hi"), "a.pdf", "application/pdf"
|
||||
)
|
||||
get = AsyncMock(return_value=(b"x", {}))
|
||||
with patch.object(conv, "s3_get_file", get):
|
||||
await conv._load_attached_file(f"conversations/{CONV_ID}/files/u.txt")
|
||||
assert resp.status_code == 200
|
||||
assert calls == [conv.extract_text, conv.extract_text]
|
||||
|
||||
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"
|
||||
@@ -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'<attached file "{safe_name}">\n{cap_text(text)}\n</attached>'
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user