From d1a57260de496325cb40c734fbf007a0b7a428f0 Mon Sep 17 00:00:00 2001 From: Zigao Wang Date: Thu, 30 Jan 2025 22:49:01 +0800 Subject: [PATCH] fix: markdown and thinking process rendering --- Ollmao/Ollmao/ContentView.swift | 285 ++++++++++++++++++-------------- 1 file changed, 163 insertions(+), 122 deletions(-) diff --git a/Ollmao/Ollmao/ContentView.swift b/Ollmao/Ollmao/ContentView.swift index fac02f5..743399f 100644 --- a/Ollmao/Ollmao/ContentView.swift +++ b/Ollmao/Ollmao/ContentView.swift @@ -112,7 +112,8 @@ struct ChatView: View { } if !viewModel.currentStreamContent.isEmpty { - StreamingMessageView(content: viewModel.currentStreamContent) + StreamingMessageView(content: viewModel.currentStreamContent, isStreaming: viewModel.isStreaming) + .id("streaming") } else if viewModel.isLoading { TypingIndicator() } @@ -183,6 +184,98 @@ struct ChatView: View { } } +struct MarkdownView: View { + let content: String + + var body: some View { + Markdown(content) + .textSelection(.enabled) + .applyCodeBlockStyle() + .frame(maxWidth: .infinity, alignment: .leading) + } +} + +private extension View { + func applyCodeBlockStyle() -> some View { + markdownBlockStyle(\.codeBlock) { configuration in + VStack(alignment: .leading, spacing: 0) { + // Language label if available + if let language = configuration.language { + Text(language) + .font(.system(size: 12)) + .foregroundColor(.secondary) + .padding(.horizontal, 12) + .padding(.vertical, 4) + } + + // Code content + ScrollView(.horizontal, showsIndicators: false) { + configuration.label + .font(.system(.body, design: .monospaced)) + .padding(12) + .frame(maxWidth: .infinity, alignment: .leading) + } + } + .background { + RoundedRectangle(cornerRadius: 8) + .fill(Color(nsColor: .textBackgroundColor)) + } + .overlay(alignment: .topTrailing) { + Button { + #if os(iOS) + UIPasteboard.general.string = configuration.content + #else + NSPasteboard.general.clearContents() + NSPasteboard.general.setString(configuration.content, forType: .string) + #endif + } label: { + Image(systemName: "doc.on.doc") + .foregroundColor(.secondary) + .padding(8) + } + .buttonStyle(.plain) + } + .overlay { + RoundedRectangle(cornerRadius: 8) + .stroke(Color(nsColor: .separatorColor), lineWidth: 1) + } + .padding(.vertical, 8) // Add vertical margin + } + } + + func applyHeadingStyles() -> some View { + self + .markdownBlockStyle(\.heading1) { config in + config.label + .foregroundColor(.primary) + .font(.system(size: 28, weight: .bold)) + .padding(.vertical, 8) + } + .markdownBlockStyle(\.heading2) { config in + config.label + .foregroundColor(.primary) + .font(.system(size: 24, weight: .bold)) + .padding(.vertical, 6) + } + .markdownBlockStyle(\.heading3) { config in + config.label + .foregroundColor(.primary) + .font(.system(size: 20, weight: .bold)) + .padding(.vertical, 4) + } + } + + func applyParagraphStyle() -> some View { + markdownBlockStyle(\.paragraph) { config in + config.label + .foregroundColor(.primary) + .font(.system(size: 16)) + .lineSpacing(4) + .padding(.vertical, 2) + } + } +} + struct MessageView: View { let message: ChatMessage @@ -225,11 +318,7 @@ struct MessageView: View { if message.content.contains("") { ThinkingStreamView(content: message.content) } else { - ScrollView { - Text(.init(message.content)) - .textSelection(.enabled) - .frame(maxWidth: .infinity, alignment: .leading) - } + MarkdownView(content: message.content) } } } @@ -239,148 +328,108 @@ struct MessageView: View { } } -struct ThinkingView: View { +struct ThinkingStreamView: View { let content: String - @State private var isThinkingExpanded = true + let isStreaming: Bool @State private var brainScale: CGFloat = 1.0 + @State private var isExpanded: Bool = true - var body: some View { - let thinkingContent = extractThinkingContent(from: content) - if content.contains("") { - VStack(alignment: .leading, spacing: 8) { - HStack { - Image(systemName: "brain.head.profile") - .foregroundColor(.purple) - .scaleEffect(brainScale) - Text(thinkingContent.isEmpty ? "No thinking process" : "Thinking Process") - .foregroundColor(.purple) - } - .padding(8) - .background(Color.purple.opacity(0.1)) - .cornerRadius(8) - .onAppear { - withAnimation(.easeInOut(duration: 0.6).repeatForever()) { - brainScale = 1.1 - } - } - - ScrollView { - Text(.init(extractFinalAnswer(from: content))) - .textSelection(.enabled) - .padding(.top, 8) - } - } - } else { - ScrollView { - Text(.init(content)) - .textSelection(.enabled) - .frame(maxWidth: .infinity, alignment: .leading) - } - } - } - - private func extractThinkingContent(from content: String) -> String { - if let start = content.range(of: "")?.upperBound, - let end = content.range(of: "")?.lowerBound { - let thinking = String(content[start.. String { - if let end = content.range(of: "")?.upperBound { - return String(content[end...]).trimmingCharacters(in: .whitespacesAndNewlines) - } - return content + init(content: String, isStreaming: Bool = false) { + self.content = content + self.isStreaming = isStreaming + _isExpanded = State(initialValue: isStreaming) } -} - -struct ThinkingStreamView: View { - let content: String - @State private var brainScale: CGFloat = 1.0 var body: some View { - let (thinkingContent, finalAnswer) = extractContent(from: content) + let thinkingContent = extractThinkingContent(from: content) VStack(alignment: .leading, spacing: 12) { - // Only show thinking process if it exists - if !thinkingContent.isEmpty { - VStack(alignment: .leading, spacing: 8) { + // Thinking process section + VStack(alignment: .leading, spacing: 8) { + Button(action: { withAnimation { isExpanded.toggle() } }) { HStack { Image(systemName: "brain.head.profile") .foregroundColor(.purple) .scaleEffect(brainScale) Text("Thinking Process") .foregroundColor(.purple) + Spacer() + Image(systemName: isExpanded ? "chevron.up" : "chevron.down") + .foregroundColor(.purple) } .padding(8) .background(Color.purple.opacity(0.1)) .cornerRadius(8) - .onAppear { - withAnimation(.easeInOut(duration: 0.6).repeatForever()) { - brainScale = 1.1 + } + .buttonStyle(.plain) + .onAppear { + withAnimation(.easeInOut(duration: 0.6).repeatForever()) { + brainScale = 1.1 + } + } + + if isExpanded { + Group { + if thinkingContent.isEmpty { + Text("No thinking process") + .foregroundColor(.secondary) + .padding() + .frame(maxWidth: .infinity, alignment: .center) + .background(Color.gray.opacity(0.1)) + .cornerRadius(8) + } else { + Text(thinkingContent) + .textSelection(.enabled) + .padding() + .background(Color.gray.opacity(0.1)) + .cornerRadius(8) + .overlay( + Rectangle() + .fill(Color.purple.opacity(0.3)) + .frame(width: 4) + .padding(.vertical, 4), + alignment: .leading + ) } } - - Text(thinkingContent) - .textSelection(.enabled) - .padding() - .background(Color.gray.opacity(0.1)) - .cornerRadius(8) - .overlay( - Rectangle() - .fill(Color.purple.opacity(0.3)) - .frame(width: 4) - .padding(.vertical, 4), - alignment: .leading - ) + .transition(AnyTransition.move(edge: .top).combined(with: .opacity)) } } - // Show final answer as normal text - if !finalAnswer.isEmpty { - Text(.init(finalAnswer)) - .textSelection(.enabled) - .frame(maxWidth: .infinity, alignment: .leading) + // Final answer section + if let finalContent = extractFinalContent(from: content), !finalContent.isEmpty { + MarkdownView(content: finalContent) } } - } - - private func extractContent(from content: String) -> (thinking: String, answer: String) { - // If we don't have a complete thinking process (no ), treat everything as thinking - guard content.contains("") else { - let cleanContent = content - .replacingOccurrences(of: "", with: "") - .trimmingCharacters(in: .whitespacesAndNewlines) - return (thinking: cleanContent, answer: "") - } - - // Extract thinking content - var thinkingContent = "" - if let start = content.range(of: "")?.upperBound, - let end = content.range(of: "")?.lowerBound { - thinkingContent = String(content[start.. String { + guard let startRange = content.range(of: "") else { return "" } - // Extract final answer - var finalAnswer = "" - if let end = content.range(of: "")?.upperBound { - finalAnswer = String(content[end...]).trimmingCharacters(in: .whitespacesAndNewlines) + let afterStartTag = content[startRange.upperBound...] + if let endRange = afterStartTag.range(of: "") { + let thinking = String(afterStartTag[.. String? { + guard let endRange = content.range(of: "") else { return nil } + return String(content[endRange.upperBound...]).trimmingCharacters(in: .whitespacesAndNewlines) } } struct StreamingMessageView: View { let content: String + let isStreaming: Bool var body: some View { HStack(alignment: .top, spacing: 12) { @@ -408,15 +457,7 @@ struct StreamingMessageView: View { .buttonStyle(.plain) } - if content.contains("") { - ThinkingStreamView(content: content) - } else { - ScrollView { - Text(.init(content)) - .textSelection(.enabled) - .frame(maxWidth: .infinity, alignment: .leading) - } - } + ThinkingStreamView(content: content, isStreaming: isStreaming) } } .padding()