From 0b0a5815a6d45dd3ffac95742df7e43d7ba23ab1 Mon Sep 17 00:00:00 2001
From: Zigao Wang <zigao.wang24@gmail.com>
Date: Wed, 5 Feb 2025 13:21:46 +0800
Subject: [PATCH] feat: update context prompt easier

---
 Ollmao/Ollmao/Services/OllamaService.swift | 47 +++++++++++-----------
 1 file changed, 24 insertions(+), 23 deletions(-)

diff --git a/Ollmao/Ollmao/Services/OllamaService.swift b/Ollmao/Ollmao/Services/OllamaService.swift
index 37be8e1..1f554f2 100644
--- a/Ollmao/Ollmao/Services/OllamaService.swift
+++ b/Ollmao/Ollmao/Services/OllamaService.swift
@@ -13,37 +13,30 @@ actor OllamaService {
 
     private struct Context {
         static let systemPrompt = """
-        You are an AI assistant in Ollmao, a modern and user-friendly Ollama GUI application developed by Zigao Wang and Thomas Wu. As Ollmao's assistant, you aim to:
-        - Provide clear, accurate, and helpful responses that reflect Ollmao's commitment to quality
-        - Maintain a friendly and approachable tone while being professional
-        - Help users understand and make the most of Ollmao's features
-        - Respect user privacy and data security
-        - Format responses in a clean, easy-to-read manner
-        - Admit when you're unsure and provide accurate information about Ollmao's capabilities
-        - Stay up-to-date with Ollama's features and limitations
-
-        Current conversation mode: chat
+        You are a helpful assistant. Be clear and direct in your responses.
         """
         
         static let contextTemplate = """
-        [INST]
-        <<SYS>>
+        [SYSTEM]
         \(systemPrompt)
-        <</SYS>>
+        [/SYSTEM]
         
         {history}
         
-        Human: {input}
-        [/INST]
+        [HUMAN]
+        {input}
+        [/HUMAN]
+        
+        [ASSISTANT]
         """
     }
 
     private func formatContext(_ messages: [ChatMessage]) -> String {
         let history = messages.map { message in
             switch message.role {
-                case .user: return "Human: \(message.content)"
-                case .assistant: return "Assistant: \(message.content)"
-                case .system: return message.content // System messages are included directly
+                case .user: return message.content
+                case .assistant: return message.content
+                case .system: return message.content
             }
         }.joined(separator: "\n\n")
         
@@ -52,12 +45,20 @@ actor OllamaService {
     }
     
     private func cleanResponse(_ response: String) -> String {
-        // Remove any "Assistant:" or "Human:" prefixes from the response
-        let cleaned = response.trimmingCharacters(in: .whitespacesAndNewlines)
-        if cleaned.lowercased().hasPrefix("assistant:") {
-            return String(cleaned.dropFirst("assistant:".count)).trimmingCharacters(in: .whitespacesAndNewlines)
+        var cleaned = response.trimmingCharacters(in: .whitespacesAndNewlines)
+        
+        // Remove role tags and their content
+        let roleTags = ["[HUMAN]", "[ASSISTANT]", "[SYSTEM]"]
+        for tag in roleTags {
+            while let range = cleaned.range(of: "\\[\(tag)\\].*?\\[\\/\(tag)\\]", options: .regularExpression) {
+                cleaned.removeSubrange(range)
+            }
+            cleaned = cleaned.replacingOccurrences(of: tag, with: "")
+            cleaned = cleaned.replacingOccurrences(of: "[\(tag)]", with: "")
+            cleaned = cleaned.replacingOccurrences(of: "[/\(tag)]", with: "")
         }
-        return cleaned
+        
+        return cleaned.trimmingCharacters(in: .whitespacesAndNewlines)
     }
     
     func generateResponse(prompt: String, messages: [ChatMessage], model: String) async throws -> AsyncThrowingStream<String, Error> {