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 starlette.responses import StreamingResponse, FileResponse
from pydantic import BaseModel
from fpdf import FPDF
import markdown
import requests
import os
import aiohttp
import json
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 constants import ERROR_MESSAGES
from typing import List
router = APIRouter()
@ -41,6 +36,59 @@ async def get_html_from_markdown(
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")
async def download_db(user=Depends(get_admin_user)):

View file

@ -42,6 +42,8 @@ xlrd
opencv-python-headless
rapidocr-onnxruntime
fpdf2
faster-whisper
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;
};
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) => {
let error = null;

View file

@ -11,6 +11,8 @@
import Dropdown from '$lib/components/common/Dropdown.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 shareHandler: Function;
@ -25,7 +27,7 @@
export let onClose: Function = () => {};
const downloadChatAsTxt = async () => {
const downloadTxt = async () => {
const _chat = chat.chat;
console.log('download', chat);
@ -40,54 +42,29 @@
saveAs(blob, `chat-${_chat.title}.txt`);
};
const downloadChatAsPdf = async () => {
const downloadPdf = async () => {
const _chat = chat.chat;
console.log('download', chat);
const doc = new jsPDF();
const blob = await downloadChatAsPDF(_chat);
// Initialize y-coordinate for text placement
let yPos = 10;
const pageHeight = doc.internal.pageSize.height;
// Create a URL for the blob
const url = window.URL.createObjectURL(blob);
// Function to check if new text exceeds the current page height
function checkAndAddNewPage() {
if (yPos > pageHeight - 10) {
doc.addPage();
yPos = 10; // Reset yPos for the new page
}
}
// Create a link element to trigger the download
const a = document.createElement('a');
a.href = url;
a.download = `chat-${_chat.title}.pdf`;
// Function to add text with specific style
function addStyledText(text, isTitle = false) {
// Set font style and size based on the parameters
doc.setFont('helvetica', isTitle ? 'bold' : 'normal');
doc.setFontSize(isTitle ? 12 : 10);
// Append the link to the body and click it programmatically
document.body.appendChild(a);
a.click();
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
const lines = doc.splitTextToSize(text, 180); // Adjust the width as needed
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`);
// Revoke the URL to release memory
window.URL.revokeObjectURL(url);
};
</script>
@ -193,7 +170,7 @@
<DropdownMenu.Item
class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md"
on:click={() => {
downloadChatAsTxt();
downloadTxt();
}}
>
<div class="flex items-center line-clamp-1">Plain text (.txt)</div>
@ -202,7 +179,7 @@
<DropdownMenu.Item
class="flex gap-2 items-center px-3 py-2 text-sm cursor-pointer dark:hover:bg-gray-850 rounded-md"
on:click={() => {
downloadChatAsPdf();
downloadPdf();
}}
>
<div class="flex items-center line-clamp-1">PDF document (.pdf)</div>