refac: pdf generation

This commit is contained in:
Timothy J. Baek 2024-04-06 01:54:59 -07:00
parent d001d7afb1
commit 81dbc65853
9 changed files with 104 additions and 51 deletions

View file

@ -1,16 +1,11 @@
from fastapi import APIRouter, UploadFile, File, BackgroundTasks from fastapi import APIRouter, UploadFile, File, Response
from fastapi import Depends, HTTPException, status from fastapi import Depends, HTTPException, status
from starlette.responses import StreamingResponse, FileResponse from starlette.responses import StreamingResponse, FileResponse
from pydantic import BaseModel from pydantic import BaseModel
from fpdf import FPDF
import markdown import markdown
import requests
import os
import aiohttp
import json
from utils.utils import get_admin_user from utils.utils import get_admin_user
@ -18,7 +13,7 @@ from utils.misc import calculate_sha256, get_gravatar_url
from config import OLLAMA_BASE_URLS, DATA_DIR, UPLOAD_DIR from config import OLLAMA_BASE_URLS, DATA_DIR, UPLOAD_DIR
from constants import ERROR_MESSAGES from constants import ERROR_MESSAGES
from typing import List
router = APIRouter() router = APIRouter()
@ -41,6 +36,59 @@ async def get_html_from_markdown(
return {"html": markdown.markdown(form_data.md)} return {"html": markdown.markdown(form_data.md)}
class ChatForm(BaseModel):
title: str
messages: List[dict]
@router.post("/pdf")
async def download_chat_as_pdf(
form_data: ChatForm,
):
pdf = FPDF()
pdf.add_page()
STATIC_DIR = "./static"
FONTS_DIR = f"{STATIC_DIR}/fonts"
pdf.add_font("NotoSans", "", f"{FONTS_DIR}/NotoSans-Regular.ttf")
pdf.add_font("NotoSans", "b", f"{FONTS_DIR}/NotoSans-Bold.ttf")
pdf.add_font("NotoSans", "i", f"{FONTS_DIR}/NotoSans-Italic.ttf")
pdf.add_font("NotoSansKR", "", f"{FONTS_DIR}/NotoSansKR-Regular.ttf")
pdf.add_font("NotoSansJP", "", f"{FONTS_DIR}/NotoSansJP-Regular.ttf")
pdf.set_font("NotoSans", size=12)
pdf.set_fallback_fonts(["NotoSansKR", "NotoSansJP"])
pdf.set_auto_page_break(auto=True, margin=15)
# Adjust the effective page width for multi_cell
effective_page_width = (
pdf.w - 2 * pdf.l_margin - 10
) # Subtracted an additional 10 for extra padding
# Add chat messages
for message in form_data.messages:
role = message["role"]
content = message["content"]
pdf.set_font("NotoSans", "B", size=12) # Bold for the role
pdf.multi_cell(effective_page_width, 10, f"{role.upper()}", 0, "L")
pdf.ln(1) # Extra space between messages
pdf.set_font("NotoSans", size=10) # Regular for content
pdf.multi_cell(effective_page_width, 10, content, 0, "L")
pdf.ln(1) # Extra space between messages
# Save the pdf with name .pdf
pdf_bytes = pdf.output()
return Response(
content=bytes(pdf_bytes),
media_type="application/pdf",
headers={"Content-Disposition": f"attachment;filename=chat.pdf"},
)
@router.get("/db/download") @router.get("/db/download")
async def download_db(user=Depends(get_admin_user)): async def download_db(user=Depends(get_admin_user)):

View file

@ -42,6 +42,8 @@ xlrd
opencv-python-headless opencv-python-headless
rapidocr-onnxruntime rapidocr-onnxruntime
fpdf2
faster-whisper faster-whisper
PyJWT PyJWT

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

Binary file not shown.

View file

@ -22,6 +22,32 @@ export const getGravatarUrl = async (email: string) => {
return res; return res;
}; };
export const downloadChatAsPDF = async (chat: object) => {
let error = null;
const blob = await fetch(`${WEBUI_API_BASE_URL}/utils/pdf`, {
method: 'POST',
headers: {
'Content-Type': 'application/json'
},
body: JSON.stringify({
title: chat.title,
messages: chat.messages
})
})
.then(async (res) => {
if (!res.ok) throw await res.json();
return res.blob();
})
.catch((err) => {
console.log(err);
error = err;
return null;
});
return blob;
};
export const getHTMLFromMarkdown = async (md: string) => { export const getHTMLFromMarkdown = async (md: string) => {
let error = null; let error = null;

View file

@ -11,6 +11,8 @@
import Dropdown from '$lib/components/common/Dropdown.svelte'; import Dropdown from '$lib/components/common/Dropdown.svelte';
import Tags from '$lib/components/common/Tags.svelte'; import Tags from '$lib/components/common/Tags.svelte';
import { WEBUI_BASE_URL } from '$lib/constants';
import { downloadChatAsPDF } from '$lib/apis/utils';
export let shareEnabled: boolean = false; export let shareEnabled: boolean = false;
export let shareHandler: Function; export let shareHandler: Function;
@ -25,7 +27,7 @@
export let onClose: Function = () => {}; export let onClose: Function = () => {};
const downloadChatAsTxt = async () => { const downloadTxt = async () => {
const _chat = chat.chat; const _chat = chat.chat;
console.log('download', chat); console.log('download', chat);
@ -40,54 +42,29 @@
saveAs(blob, `chat-${_chat.title}.txt`); saveAs(blob, `chat-${_chat.title}.txt`);
}; };
const downloadChatAsPdf = async () => { const downloadPdf = async () => {
const _chat = chat.chat; const _chat = chat.chat;
console.log('download', chat); console.log('download', chat);
const doc = new jsPDF(); const blob = await downloadChatAsPDF(_chat);
// Initialize y-coordinate for text placement // Create a URL for the blob
let yPos = 10; const url = window.URL.createObjectURL(blob);
const pageHeight = doc.internal.pageSize.height;
// Function to check if new text exceeds the current page height // Create a link element to trigger the download
function checkAndAddNewPage() { const a = document.createElement('a');
if (yPos > pageHeight - 10) { a.href = url;
doc.addPage(); a.download = `chat-${_chat.title}.pdf`;
yPos = 10; // Reset yPos for the new page
}
}
// Function to add text with specific style // Append the link to the body and click it programmatically
function addStyledText(text, isTitle = false) { document.body.appendChild(a);
// Set font style and size based on the parameters a.click();
doc.setFont('helvetica', isTitle ? 'bold' : 'normal');
doc.setFontSize(isTitle ? 12 : 10);
const textMargin = 7; // Remove the link from the body
document.body.removeChild(a);
// Split text into lines to ensure it fits within the page width // Revoke the URL to release memory
const lines = doc.splitTextToSize(text, 180); // Adjust the width as needed window.URL.revokeObjectURL(url);
lines.forEach((line) => {
checkAndAddNewPage(); // Check if we need a new page before adding more text
doc.text(line, 10, yPos);
yPos += textMargin; // Increment yPos for the next line
});
// Add extra space after a block of text
yPos += 2;
}
_chat.messages.forEach((message, i) => {
// Add user text in bold
doc.setFont('helvetica', 'normal', 'bold');
addStyledText(message.role.toUpperCase(), { isTitle: true });
addStyledText(message.content);
});
doc.save(`chat-${_chat.title}.pdf`);
}; };
</script> </script>
@ -193,7 +170,7 @@
<DropdownMenu.Item <DropdownMenu.Item
class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md" class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md"
on:click={() => { on:click={() => {
downloadChatAsTxt(); downloadTxt();
}} }}
> >
<div class="flex items-center line-clamp-1">Plain text (.txt)</div> <div class="flex items-center line-clamp-1">Plain text (.txt)</div>
@ -202,7 +179,7 @@
<DropdownMenu.Item <DropdownMenu.Item
class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md" class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md"
on:click={() => { on:click={() => {
downloadChatAsPdf(); downloadPdf();
}} }}
> >
<div class="flex items-center line-clamp-1">PDF document (.pdf)</div> <div class="flex items-center line-clamp-1">PDF document (.pdf)</div>