Merge pull request #37 from Promptly-Technologies-LLC/11-handle-message-annotations-for-rendering-citations-to-files-referenced-by-the-assistant-run

Convert file citation text deltas to download links

Author: chriscarrollsmith

.gitignore

@@ -9,4 +9,5 @@ __pycache__
 .cursorrules
 .repomix-output.txt
 repomix-output.txt
-artifacts/
+artifacts/
+uploads/

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "openai-assistants-python-quickstart"
-version = "0.1.0"
+version = "1.0.0"
 description = "A quickstart template using the OpenAI Assistants API with Python, FastAPI, Jinja2, and HTMX"
 readme = "README.md"
 requires-python = ">=3.12"

routers/chat.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from datetime import datetime
 from typing import AsyncGenerator, Optional, Union
 from fastapi.templating import Jinja2Templates
@@ -19,6 +20,7 @@
 
 import json
 
+from routers import files as files_router
 from utils.custom_functions import get_weather, post_tool_outputs
 from utils.sse import sse_format
 from utils.streaming import AssistantStreamMetadata
@@ -121,6 +123,28 @@ async def handle_assistant_stream(
                     if isinstance(content, TextDeltaBlock) and content.text and content.text.value:
                         step_id = event.data.id
                         text_value = content.text.value
+                        annotations = content.text.annotations
+                        
+                        # Check for file citation annotations
+                        if annotations:
+                            logger.debug(f"Annotation found: {annotations}")
+                            for annotation in annotations:
+                                if annotation.type == 'file_citation' and hasattr(annotation, 'file_citation') and annotation.file_citation:
+                                    match = re.search(r'【.*?†(.*?)】', text_value)
+                                    if match:
+                                        file_name = match.group(1)
+                                        # Manually construct the URL
+                                        file_url = f'/assistants/{assistant_id}/files/{file_name}'
+                                        text_value = f'[†]({file_url})'
+                                        logger.debug(f"Replacing citation with link: {text_value}")
+                                    else:
+                                        # Handle error: pattern not found in the text
+                                        logger.warning(f"Could not extract filename from citation text: {text_value}")
+                                        file_name = None # Indicate failure
+                                    
+                                    # Assuming one citation per delta for now
+                                    break 
+                        
                         sse_data = f'<span hx-swap-oob="beforeend:#step-{step_id}">{text_value}</span>'
                         yield sse_format(
                             "textDelta",
@@ -205,7 +229,7 @@ async def handle_assistant_stream(
     async def event_generator() -> AsyncGenerator[str, None]:
         """
         Main generator for SSE events. We call our helper function to handle the assistant
-        stream, and if the assistant requests a tool call, we do it and then re-stream the stream.
+        stream, and if the assistant requests a tool call, we execute it and then re-stream the stream.
         """
         step_id: str = ""
         stream_manager: AsyncAssistantStreamManager[AsyncAssistantEventHandler] = client.beta.threads.runs.stream(

routers/files.py

@@ -5,10 +5,10 @@
 from fastapi import (
     APIRouter, Request, UploadFile, File, HTTPException, Depends, Path, Form
 )
-from fastapi.responses import HTMLResponse, StreamingResponse
+from fastapi.responses import HTMLResponse, StreamingResponse, FileResponse
 from fastapi.templating import Jinja2Templates
 from openai import AsyncOpenAI
-from utils.files import get_or_create_vector_store, get_files_for_vector_store
+from utils.files import get_or_create_vector_store, get_files_for_vector_store, store_file, retrieve_file, delete_local_file
 from utils.streaming import stream_file_content
 
 logger = logging.getLogger("uvicorn.error")
@@ -69,18 +69,38 @@ async def upload_file(
         )
 
     error_message = None
+    file_content = None
     try:
-        # Upload the file to OpenAI
+        # 1. Read the file content first
+        file_content = await file.read()
+        if not file.filename:
+            raise ValueError("File has no filename")
+        if not file_content:
+            raise ValueError("File content is empty")
+
+        # 2. Upload the file content to OpenAI
         openai_file = await client.files.create(
-            file=(file.filename, file.file),
+            file=(file.filename, file_content),
             purpose=purpose
         )
 
-        # Add the uploaded file to the vector store
+        # 3. Add the uploaded file to the vector store
         await client.vector_stores.files.create(
             vector_store_id=vector_store_id,
             file_id=openai_file.id
         )
+        logger.info(f"File {file.filename} uploaded to OpenAI and added to vector store.")
+
+        # 4. Store the file locally using the same content
+        try:
+            store_file(file.filename, file_content)
+        except Exception as e:
+            logger.error(f"Error storing file {file.filename} locally: {e}")
+            error_message = f"Error storing file locally"
+
+    except ValueError as ve:
+        logger.error(f"File validation error for assistant {assistant_id}: {ve}")
+        error_message = str(ve)
     except Exception as e:
         logger.error(f"Error uploading file for assistant {assistant_id}: {e}")
         error_message = f"Error uploading file for assistant"
@@ -121,6 +141,26 @@ async def delete_file(
     try:
         vector_store_id = await get_or_create_vector_store(assistant_id, client)
 
+        # Retrieve filename before attempting deletions
+        file_to_delete_name = None
+        try:
+            retrieved_file = await client.files.retrieve(file_id)
+            if retrieved_file and retrieved_file.filename:
+                file_to_delete_name = retrieved_file.filename
+                logger.info(f"Retrieved filename '{retrieved_file.filename}' for deletion.")
+            else:
+                logger.warning(f"Could not retrieve filename for file_id {file_id}")
+        except Exception as retrieve_err:
+            logger.error(f"Error retrieving file object {file_id} for filename: {retrieve_err}")
+
+        # Attempt to delete stored file if filename was found
+        if file_to_delete_name:
+            try:
+                delete_local_file(file_to_delete_name)
+            except Exception as local_delete_err:
+                # Log error but continue with OpenAI/VS deletion
+                logger.error(f"Error deleting local file {file_to_delete_name}: {local_delete_err}")
+
         # 1. Delete the file association from the vector store
         try:
             deleted_vs_file = await client.vector_stores.files.delete(
@@ -175,9 +215,20 @@ async def delete_file(
 
 # --- Streaming file content routes ---
 
-@router.get("/{file_id}")
+@router.get("/{file_name}")
 async def download_assistant_file(
-    file_id: str = Path(..., description="The ID of the file to retrieve"),
+    assistant_id: str = Path(..., description="The ID of the Assistant"),
+    file_name: str = Path(..., description="The name of the file to retrieve")
+) -> FileResponse:
+    """Serves an assistant file stored locally in the uploads directory."""
+    return retrieve_file(file_name)
+
+
+# This endpoint retrieves files uploaded TO openai (e.g., code interpreter output)
+# Keep it separate for clarity
+@router.get("/{file_id}/openai_content")
+async def download_openai_file(
+    file_id: str = Path(..., description="The ID of the file stored in OpenAI"),
     client: AsyncOpenAI = Depends(lambda: AsyncOpenAI())
 ) -> StreamingResponse:
     try:
@@ -187,12 +238,14 @@ async def download_assistant_file(
         if not hasattr(file_content, 'content'):
             raise HTTPException(status_code=500, detail="File content not available")
 
+        # Use stream_file_content helper
         return StreamingResponse(
-            stream_file_content(file_content.content),
-            headers={"Content-Disposition": f'attachment; filename=\"{file.filename or "download"}\"'}
+            stream_file_content(file_content.content), # Assuming stream_file_content handles bytes
+            headers={"Content-Disposition": f'attachment; filename="{file.filename or file_id}"'} # Use file_id as fallback filename
         )
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Error downloading file {file_id} from OpenAI: {e}")
+        raise HTTPException(status_code=500, detail=f"Error downloading file from OpenAI: {str(e)}")
 
 
 @router.get("/{file_id}/content")

utils/files.py

@@ -1,9 +1,12 @@
 import logging
-from typing import List, Dict, Any
-from fastapi import Depends
+import os
+from typing import List, Dict
+from fastapi import Depends, HTTPException
+from fastapi.responses import FileResponse
 from openai import AsyncOpenAI
 from openai.types.file_object import FileObject
 from openai.types.vector_stores.vector_store_file import VectorStoreFile
+
 logger = logging.getLogger("uvicorn.error")
 
 # Helper function to get or create a vector store
@@ -60,4 +63,61 @@ async def get_files_for_vector_store(vector_store_id: str, client: AsyncOpenAI)
     except Exception as e:
         logger.error(f"Error listing files for vector store {vector_store_id}: {e}")
         # Return empty list or re-raise depending on desired behavior
-        return []
+        return []
+
+
+def store_file(file_name: str, file_content: bytes):
+    """Store a file in the uploads directory."""
+    upload_dir = "uploads"
+    os.makedirs(upload_dir, exist_ok=True) # Create directory if it doesn't exist
+    file_path = os.path.join(upload_dir, file_name)
+    try:
+        with open(file_path, "wb") as f:
+            f.write(file_content)
+        logger.info(f"File stored successfully: {file_path}")
+    except IOError as e:
+        logger.error(f"Error writing file {file_path}: {e}")
+        # Re-raise a more specific exception or handle as needed
+        # For now, letting the caller's try/except handle it might be okay
+        raise
+
+
+def retrieve_file(file_name: str):
+    """Retrieve a file from the uploads directory."""
+    upload_dir = "uploads"
+    file_path = os.path.join(upload_dir, file_name)
+
+    if not os.path.exists(file_path):
+        logger.error(f"File not found: {file_path}")
+        raise HTTPException(status_code=404, detail="File not found")
+
+    # Basic security check: prevent path traversal
+    if not os.path.abspath(file_path).startswith(os.path.abspath(upload_dir)):
+        logger.error(f"Attempted path traversal: {file_path}")
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    try:
+        return FileResponse(path=file_path, filename=file_name, media_type='application/octet-stream')
+    except Exception as e:
+        logger.error(f"Error serving file {file_path}: {e}")
+        raise HTTPException(status_code=500, detail="Error serving file")
+
+
+def delete_local_file(file_name: str):
+    """Delete a file from the local uploads directory."""
+    upload_dir = "uploads"
+    file_path = os.path.join(upload_dir, file_name)
+
+    # Basic security check
+    if not os.path.abspath(file_path).startswith(os.path.abspath(upload_dir)):
+        logger.error(f"Attempted path traversal during delete: {file_path}")
+        raise HTTPException(status_code=403, detail="Forbidden")
+
+    try:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            logger.info(f"Successfully deleted local file: {file_path}")
+        else:
+            logger.warning(f"Attempted to delete local file, but it was not found: {file_path}")
+    except OSError as e:
+        logger.error(f"Error deleting local file {file_path}: {e}")