Files
MLXServer/MLXServer/Utilities/ChatExporter.swift

296 lines
11 KiB
Swift

import AppKit
import Foundation
import SwiftUI
import UniformTypeIdentifiers
extension UTType {
static var markdownText: UTType {
UTType(filenameExtension: "md") ?? .plainText
}
}
/// FileDocument for exporting chat as Markdown (.md) or RTF.
struct ChatExportDocument: FileDocument {
static var readableContentTypes: [UTType] { [.rtf, .markdownText] }
static var writableContentTypes: [UTType] { [.rtf, .markdownText] }
let messages: [ChatMessage]
let modelName: String?
init(messages: [ChatMessage], modelName: String?) {
self.messages = messages
self.modelName = modelName
}
init(configuration: ReadConfiguration) throws {
self.messages = []
self.modelName = nil
}
func fileWrapper(configuration: WriteConfiguration) throws -> FileWrapper {
if configuration.contentType == .rtf {
guard let data = ChatExporter.exportRTF(messages: messages, modelName: modelName) else {
throw CocoaError(.fileWriteUnknown)
}
return FileWrapper(regularFileWithContents: data)
}
let md = ChatExporter.exportMarkdown(messages: messages, modelName: modelName)
return FileWrapper(regularFileWithContents: Data(md.utf8))
}
}
/// Exports a chat conversation to Markdown or RTF (Pages-compatible) format.
enum ChatExporter {
// MARK: - Markdown export
static func exportMarkdown(messages: [ChatMessage], modelName: String?) -> String {
var lines: [String] = []
// Header
lines.append("# Chat Session")
if let modelName {
lines.append("**Model:** \(modelName)")
}
let formatter = DateFormatter()
formatter.dateStyle = .long
formatter.timeStyle = .short
if let first = messages.first {
lines.append("**Date:** \(formatter.string(from: first.timestamp))")
}
lines.append("")
lines.append("---")
lines.append("")
for message in messages {
guard message.role != .system else { continue }
if message.role == .user {
// User messages as blockquotes
lines.append("**You:**")
lines.append("")
for line in message.content.components(separatedBy: "\n") {
lines.append("> \(line)")
}
} else {
// Assistant messages: carry over original markdown
lines.append("**Assistant:**")
lines.append("")
lines.append(message.content)
}
lines.append("")
lines.append("---")
lines.append("")
}
return lines.joined(separator: "\n")
}
// MARK: - RTF export
static func exportRTF(messages: [ChatMessage], modelName: String?) -> Data? {
let doc = NSMutableAttributedString()
let bodyFont = NSFont.systemFont(ofSize: 13)
let bodyBoldFont = NSFont.boldSystemFont(ofSize: 13)
let titleFont = NSFont.boldSystemFont(ofSize: 20)
let metaFont = NSFont.systemFont(ofSize: 11)
let codeFont = NSFont.monospacedSystemFont(ofSize: 12, weight: .regular)
let bodyParagraph = NSMutableParagraphStyle()
bodyParagraph.paragraphSpacing = 8
bodyParagraph.lineSpacing = 2
let userParagraph = NSMutableParagraphStyle()
userParagraph.paragraphSpacing = 8
userParagraph.lineSpacing = 2
userParagraph.headIndent = 20
userParagraph.firstLineHeadIndent = 20
// Title
doc.append(NSAttributedString(
string: "Chat Session\n",
attributes: [.font: titleFont, .paragraphStyle: bodyParagraph]
))
// Metadata
let formatter = DateFormatter()
formatter.dateStyle = .long
formatter.timeStyle = .short
var metaText = ""
if let modelName { metaText += "Model: \(modelName) " }
if let first = messages.first {
metaText += "Date: \(formatter.string(from: first.timestamp))"
}
if !metaText.isEmpty {
doc.append(NSAttributedString(
string: metaText + "\n\n",
attributes: [.font: metaFont, .foregroundColor: NSColor.secondaryLabelColor]
))
}
for message in messages {
guard message.role != .system else { continue }
if message.role == .user {
doc.append(NSAttributedString(
string: "You\n",
attributes: [
.font: bodyBoldFont,
.foregroundColor: NSColor.systemBlue,
]
))
doc.append(NSAttributedString(
string: message.content + "\n\n",
attributes: [
.font: bodyFont,
.paragraphStyle: userParagraph,
.foregroundColor: NSColor.labelColor,
]
))
} else {
doc.append(NSAttributedString(
string: "Assistant\n",
attributes: [
.font: bodyBoldFont,
.foregroundColor: NSColor.labelColor,
]
))
let rendered = renderMarkdown(message.content, bodyFont: bodyFont, codeFont: codeFont, paragraph: bodyParagraph)
doc.append(rendered)
doc.append(NSAttributedString(string: "\n\n"))
}
doc.append(NSAttributedString(
string: "\n",
attributes: [
.strikethroughStyle: NSUnderlineStyle.single.rawValue,
.strikethroughColor: NSColor.separatorColor,
.font: NSFont.systemFont(ofSize: 4),
]
))
}
return doc.rtf(from: NSRange(location: 0, length: doc.length), documentAttributes: [
.documentType: NSAttributedString.DocumentType.rtf,
])
}
// MARK: - Markdown NSAttributedString (basic)
private static func renderMarkdown(
_ text: String,
bodyFont: NSFont,
codeFont: NSFont,
paragraph: NSParagraphStyle
) -> NSAttributedString {
let result = NSMutableAttributedString()
let lines = text.components(separatedBy: "\n")
var inCodeBlock = false
var codeBlockLines: [String] = []
for line in lines {
if line.hasPrefix("```") {
if inCodeBlock {
let code = codeBlockLines.joined(separator: "\n")
let codePara = NSMutableParagraphStyle()
codePara.paragraphSpacing = 4
codePara.headIndent = 12
codePara.firstLineHeadIndent = 12
result.append(NSAttributedString(
string: code + "\n",
attributes: [
.font: codeFont,
.foregroundColor: NSColor.secondaryLabelColor,
.backgroundColor: NSColor.quaternaryLabelColor,
.paragraphStyle: codePara,
]
))
codeBlockLines = []
inCodeBlock = false
} else {
inCodeBlock = true
}
continue
}
if inCodeBlock {
codeBlockLines.append(line)
continue
}
if line.hasPrefix("### ") {
result.append(NSAttributedString(
string: String(line.dropFirst(4)) + "\n",
attributes: [.font: NSFont.boldSystemFont(ofSize: 14), .paragraphStyle: paragraph]
))
} else if line.hasPrefix("## ") {
result.append(NSAttributedString(
string: String(line.dropFirst(3)) + "\n",
attributes: [.font: NSFont.boldSystemFont(ofSize: 15), .paragraphStyle: paragraph]
))
} else if line.hasPrefix("# ") {
result.append(NSAttributedString(
string: String(line.dropFirst(2)) + "\n",
attributes: [.font: NSFont.boldSystemFont(ofSize: 17), .paragraphStyle: paragraph]
))
} else {
let styled = applyInlineFormatting(line, bodyFont: bodyFont, codeFont: codeFont)
result.append(styled)
result.append(NSAttributedString(string: "\n", attributes: [.font: bodyFont]))
}
}
return result
}
private static func applyInlineFormatting(
_ text: String,
bodyFont: NSFont,
codeFont: NSFont
) -> NSAttributedString {
let result = NSMutableAttributedString()
var remaining = text[text.startIndex...]
while !remaining.isEmpty {
if remaining.hasPrefix("`"), let end = remaining.dropFirst().firstIndex(of: "`") {
let code = String(remaining[remaining.index(after: remaining.startIndex)..<end])
result.append(NSAttributedString(
string: code,
attributes: [
.font: codeFont,
.foregroundColor: NSColor.secondaryLabelColor,
.backgroundColor: NSColor.quaternaryLabelColor,
]
))
remaining = remaining[remaining.index(after: end)...]
} else if remaining.hasPrefix("**"), let end = remaining.dropFirst(2).range(of: "**") {
let bold = String(remaining[remaining.index(remaining.startIndex, offsetBy: 2)..<end.lowerBound])
result.append(NSAttributedString(
string: bold,
attributes: [.font: NSFont.boldSystemFont(ofSize: bodyFont.pointSize)]
))
remaining = remaining[end.upperBound...]
} else if remaining.hasPrefix("*"), let end = remaining.dropFirst().firstIndex(of: "*") {
let italic = String(remaining[remaining.index(after: remaining.startIndex)..<end])
result.append(NSAttributedString(
string: italic,
attributes: [.font: NSFontManager.shared.convert(bodyFont, toHaveTrait: .italicFontMask)]
))
remaining = remaining[remaining.index(after: end)...]
} else {
let ch = remaining[remaining.startIndex]
result.append(NSAttributedString(
string: String(ch),
attributes: [.font: bodyFont]
))
remaining = remaining[remaining.index(after: remaining.startIndex)...]
}
}
return result
}
}