Add WebChatContextCollector

Author: intitni

Core/Package.swift

@@ -184,11 +184,16 @@ let package = Package(
                 "MathChatPlugin",
                 "SearchChatPlugin",
                 "ShortcutChatPlugin",
-
+                
+                // context collectors
+                "WebChatContextCollector",
+                
+                .product(name: "Parsing", package: "swift-parsing"),
                 .product(name: "OpenAIService", package: "Tool"),
                 .product(name: "Preferences", package: "Tool"),
             ]
         ),
+        .testTarget(name: "ChatServiceTests", dependencies: ["ChatService"]),
         .target(
             name: "ChatPlugin",
             dependencies: [
@@ -313,6 +318,7 @@ let package = Package(
                 "ChatPlugin",
                 .product(name: "OpenAIService", package: "Tool"),
                 .product(name: "LangChain", package: "Tool"),
+                .product(name: "ExternalServices", package: "Tool"),
                 .product(name: "PythonKit", package: "PythonKit"),
             ],
             path: "Sources/ChatPlugins/SearchChatPlugin"
@@ -327,6 +333,19 @@ let package = Package(
             ],
             path: "Sources/ChatPlugins/ShortcutChatPlugin"
         ),
+        
+        // MAKR: - Chat Context Collector
+        
+        .target(
+            name: "WebChatContextCollector",
+            dependencies: [
+                "ChatContextCollector",
+                .product(name: "OpenAIService", package: "Tool"),
+                .product(name: "ExternalServices", package: "Tool"),
+                .product(name: "Preferences", package: "Tool"),
+            ],
+            path: "Sources/ChatContextCollectors/WebChatContextCollector"
+        )
     ]
 )
 

Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift

@@ -0,0 +1,145 @@
+import BingSearchService
+import ChatContextCollector
+import Foundation
+import OpenAIService
+import Preferences
+import SuggestionModel
+
+public struct WebChatContextCollector: ChatContextCollector {
+    public init() {}
+
+    public func generateContext(
+        history: [ChatMessage],
+        scopes: Set<String>,
+        content: String
+    ) -> ChatContext? {
+        guard scopes.contains("web") else { return nil }
+        return .init(
+            systemPrompt: "You prefer to answer questions with latest latest on the internet.",
+            functions: [
+                SearchFunction(),
+            ]
+        )
+    }
+
+    struct SearchFunction: ChatGPTFunction {
+        static let dateFormatter = {
+            let it = DateFormatter()
+            it.dateFormat = "yyyy-MM-dd"
+            return it
+        }()
+
+        struct Arguments: Codable {
+            var query: String
+            var freshness: String?
+        }
+
+        struct Result: ChatGPTFunctionResult {
+            var result: BingSearchResult
+
+            var botReadableContent: String {
+                result.webPages.value.enumerated().map {
+                    let (index, page) = $0
+                    return """
+                    \(index + 1). \(page.name) \(page.url)
+                    \(page.snippet)
+                    """
+                }.joined(separator: "\n")
+            }
+        }
+
+        var name: String {
+            "searchWeb"
+        }
+
+        var description: String {
+            "Useful for when you need to answer questions about latest information."
+        }
+
+        var argumentSchema: JSONSchemaValue {
+            let today = Self.dateFormatter.string(from: Date())
+            return [
+                .type: "object",
+                .properties: [
+                    "query": [
+                        .type: "string",
+                        .description: "the search query",
+                    ],
+                    "freshness": [
+                        .type: "string",
+                        .description: .string(
+                            "limit the search result to a specific range, use only when user ask the question about current events. Today is \(today). Format: yyyy-MM-dd..yyyy-MM-dd"
+                        ),
+                        .examples: ["1919-10-20..1988-10-20"],
+                    ],
+                ],
+                .required: ["query"],
+            ]
+        }
+
+        func message(at phase: ChatGPTFunctionCallPhase) -> String {
+            func parseArgument(_ string: String) throws -> Arguments {
+                try JSONDecoder().decode(Arguments.self, from: string.data(using: .utf8) ?? Data())
+            }
+
+            switch phase {
+            case .detected:
+                return "Searching.."
+            case let .processing(argumentsJsonString):
+                do {
+                    let arguments = try parseArgument(argumentsJsonString)
+                    return "Searching \(arguments.query)"
+                } catch {
+                    return "Searching.."
+                }
+            case let .ended(argumentsJsonString, result):
+                do {
+                    let arguments = try parseArgument(argumentsJsonString)
+                    if let result = result as? Result {
+                        return """
+                        Finish searching \(arguments.query)
+                        \(
+                            result.result.webPages.value
+                                .map { "- [\($0.name)](\($0.url))" }
+                                .joined(separator: "\n")
+                        )
+                        """
+                    }
+                    return "Finish searching \(arguments.query)"
+                } catch {
+                    return "Finish searching"
+                }
+            case let .error(argumentsJsonString, _):
+                do {
+                    let arguments = try parseArgument(argumentsJsonString)
+                    return "Failed searching \(arguments.query)"
+                } catch {
+                    return "Failed searching"
+                }
+            }
+        }
+
+        func call(arguments: Arguments) async throws -> Result {
+            let bingSearch = BingSearchService(
+                subscriptionKey: UserDefaults.shared.value(for: \.bingSearchSubscriptionKey),
+                searchURL: UserDefaults.shared.value(for: \.bingSearchEndpoint)
+            )
+            let result = try await bingSearch.search(
+                query: arguments.query,
+                numberOfResult: UserDefaults.shared.value(for: \.chatGPTMaxToken) > 5000 ? 5 : 3,
+                freshness: arguments.freshness
+            )
+
+            let content = result.webPages.value.enumerated().map {
+                let (index, page) = $0
+                return """
+                \(index + 1). \(page.name) \(page.url)
+                \(page.snippet)
+                """
+            }.joined(separator: "\n")
+
+            return .init(result: result)
+        }
+    }
+}
+

Core/Sources/ChatService/AllContextCollector.swift

@@ -0,0 +1,8 @@
+import ChatContextCollector
+import WebChatContextCollector
+
+let allContextCollectors: [any ChatContextCollector] = [
+    ActiveDocumentChatContextCollector(),
+    WebChatContextCollector(),
+]
+

Core/Tests/ChatServiceTests/ParseScopesTests.swift

@@ -0,0 +1,39 @@
+import XCTest
+
+@testable import ChatService
+
+final class ParseScopesTests: XCTestCase {
+    let parse = DynamicContextController.parseScopes
+    
+    func test_parse_single_scope() async throws {
+        var prompt = "@web hello"
+        let scopes = parse(&prompt)
+        XCTAssertEqual(scopes, ["web"])
+        XCTAssertEqual(prompt, "hello")
+    }
+    
+    func test_parse_multiple_spaces() async throws {
+        var prompt = "@web                hello"
+        let scopes = parse(&prompt)
+        XCTAssertEqual(scopes, ["web"])
+        XCTAssertEqual(prompt, "hello")
+    }
+    
+    func test_parse_no_prefix_at_mark() async throws {
+        var prompt = "  @web                hello"
+        let scopes = parse(&prompt)
+        XCTAssertEqual(scopes, [])
+        XCTAssertEqual(prompt, prompt)
+    }
+    
+    func test_parse_multiple_scopes() async throws {
+        var prompt = "@web+file+selection hello"
+        let scopes = parse(&prompt)
+        XCTAssertEqual(scopes, ["web", "file", "selection"])
+        XCTAssertEqual(prompt, "hello")
+    }
+}
+
+
+
+

TestPlan.xctestplan

@@ -77,6 +77,13 @@
         "identifier" : "OpenAIServiceTests",
         "name" : "OpenAIServiceTests"
       }
+    },
+    {
+      "target" : {
+        "containerPath" : "container:Core",
+        "identifier" : "ChatServiceTests",
+        "name" : "ChatServiceTests"
+      }
     }
   ],
   "version" : 1