Skip to content

Commit 1224ccd

Browse files
intitniintitni
committed
Add WebChatContextCollector
1 parent 6a52b24 commit 1224ccd

File tree

5 files changed

+219
-1
lines changed

5 files changed

+219
-1
lines changed

‎Core/Package.swift‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,11 +184,16 @@ let package = Package(
184184
"MathChatPlugin",
185185
"SearchChatPlugin",
186186
"ShortcutChatPlugin",
187-
187+
188+
// context collectors
189+
"WebChatContextCollector",
190+
191+
.product(name: "Parsing", package: "swift-parsing"),
188192
.product(name: "OpenAIService", package: "Tool"),
189193
.product(name: "Preferences", package: "Tool"),
190194
]
191195
),
196+
.testTarget(name: "ChatServiceTests", dependencies: ["ChatService"]),
192197
.target(
193198
name: "ChatPlugin",
194199
dependencies: [
@@ -313,6 +318,7 @@ let package = Package(
313318
"ChatPlugin",
314319
.product(name: "OpenAIService", package: "Tool"),
315320
.product(name: "LangChain", package: "Tool"),
321+
.product(name: "ExternalServices", package: "Tool"),
316322
.product(name: "PythonKit", package: "PythonKit"),
317323
],
318324
path: "Sources/ChatPlugins/SearchChatPlugin"
@@ -327,6 +333,19 @@ let package = Package(
327333
],
328334
path: "Sources/ChatPlugins/ShortcutChatPlugin"
329335
),
336+
337+
// MAKR: - Chat Context Collector
338+
339+
.target(
340+
name: "WebChatContextCollector",
341+
dependencies: [
342+
"ChatContextCollector",
343+
.product(name: "OpenAIService", package: "Tool"),
344+
.product(name: "ExternalServices", package: "Tool"),
345+
.product(name: "Preferences", package: "Tool"),
346+
],
347+
path: "Sources/ChatContextCollectors/WebChatContextCollector"
348+
)
330349
]
331350
)
332351

‎Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,145 @@
1+
import BingSearchService
2+
import ChatContextCollector
3+
import Foundation
4+
import OpenAIService
5+
import Preferences
6+
import SuggestionModel
7+
8+
public struct WebChatContextCollector: ChatContextCollector {
9+
public init() {}
10+
11+
public func generateContext(
12+
history: [ChatMessage],
13+
scopes: Set<String>,
14+
content: String
15+
) -> ChatContext? {
16+
guard scopes.contains("web") else { return nil }
17+
return .init(
18+
systemPrompt: "You prefer to answer questions with latest latest on the internet.",
19+
functions: [
20+
SearchFunction(),
21+
]
22+
)
23+
}
24+
25+
struct SearchFunction: ChatGPTFunction {
26+
static let dateFormatter = {
27+
let it = DateFormatter()
28+
it.dateFormat = "yyyy-MM-dd"
29+
return it
30+
}()
31+
32+
struct Arguments: Codable {
33+
var query: String
34+
var freshness: String?
35+
}
36+
37+
struct Result: ChatGPTFunctionResult {
38+
var result: BingSearchResult
39+
40+
var botReadableContent: String {
41+
result.webPages.value.enumerated().map {
42+
let (index, page) = $0
43+
return """
44+
\(index + 1). \(page.name) \(page.url)
45+
\(page.snippet)
46+
"""
47+
}.joined(separator: "\n")
48+
}
49+
}
50+
51+
var name: String {
52+
"searchWeb"
53+
}
54+
55+
var description: String {
56+
"Useful for when you need to answer questions about latest information."
57+
}
58+
59+
var argumentSchema: JSONSchemaValue {
60+
let today = Self.dateFormatter.string(from: Date())
61+
return [
62+
.type: "object",
63+
.properties: [
64+
"query": [
65+
.type: "string",
66+
.description: "the search query",
67+
],
68+
"freshness": [
69+
.type: "string",
70+
.description: .string(
71+
"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"
72+
),
73+
.examples: ["1919-10-20..1988-10-20"],
74+
],
75+
],
76+
.required: ["query"],
77+
]
78+
}
79+
80+
func message(at phase: ChatGPTFunctionCallPhase) -> String {
81+
func parseArgument(_ string: String) throws -> Arguments {
82+
try JSONDecoder().decode(Arguments.self, from: string.data(using: .utf8) ?? Data())
83+
}
84+
85+
switch phase {
86+
case .detected:
87+
return "Searching.."
88+
case let .processing(argumentsJsonString):
89+
do {
90+
let arguments = try parseArgument(argumentsJsonString)
91+
return "Searching \(arguments.query)"
92+
} catch {
93+
return "Searching.."
94+
}
95+
case let .ended(argumentsJsonString, result):
96+
do {
97+
let arguments = try parseArgument(argumentsJsonString)
98+
if let result = result as? Result {
99+
return """
100+
Finish searching \(arguments.query)
101+
\(
102+
result.result.webPages.value
103+
.map { "- [\($0.name)](\($0.url))" }
104+
.joined(separator: "\n")
105+
)
106+
"""
107+
}
108+
return "Finish searching \(arguments.query)"
109+
} catch {
110+
return "Finish searching"
111+
}
112+
case let .error(argumentsJsonString, _):
113+
do {
114+
let arguments = try parseArgument(argumentsJsonString)
115+
return "Failed searching \(arguments.query)"
116+
} catch {
117+
return "Failed searching"
118+
}
119+
}
120+
}
121+
122+
func call(arguments: Arguments) async throws -> Result {
123+
let bingSearch = BingSearchService(
124+
subscriptionKey: UserDefaults.shared.value(for: \.bingSearchSubscriptionKey),
125+
searchURL: UserDefaults.shared.value(for: \.bingSearchEndpoint)
126+
)
127+
let result = try await bingSearch.search(
128+
query: arguments.query,
129+
numberOfResult: UserDefaults.shared.value(for: \.chatGPTMaxToken) > 5000 ? 5 : 3,
130+
freshness: arguments.freshness
131+
)
132+
133+
let content = result.webPages.value.enumerated().map {
134+
let (index, page) = $0
135+
return """
136+
\(index + 1). \(page.name) \(page.url)
137+
\(page.snippet)
138+
"""
139+
}.joined(separator: "\n")
140+
141+
return .init(result: result)
142+
}
143+
}
144+
}
145+

‎Core/Sources/ChatService/AllContextCollector.swift‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import ChatContextCollector
2+
import WebChatContextCollector
3+
4+
let allContextCollectors: [any ChatContextCollector] = [
5+
ActiveDocumentChatContextCollector(),
6+
WebChatContextCollector(),
7+
]
8+

‎Core/Tests/ChatServiceTests/ParseScopesTests.swift‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
import XCTest
2+
3+
@testable import ChatService
4+
5+
final class ParseScopesTests: XCTestCase {
6+
let parse = DynamicContextController.parseScopes
7+
8+
func test_parse_single_scope() async throws {
9+
var prompt = "@web hello"
10+
let scopes = parse(&prompt)
11+
XCTAssertEqual(scopes, ["web"])
12+
XCTAssertEqual(prompt, "hello")
13+
}
14+
15+
func test_parse_multiple_spaces() async throws {
16+
var prompt = "@web hello"
17+
let scopes = parse(&prompt)
18+
XCTAssertEqual(scopes, ["web"])
19+
XCTAssertEqual(prompt, "hello")
20+
}
21+
22+
func test_parse_no_prefix_at_mark() async throws {
23+
var prompt = " @web hello"
24+
let scopes = parse(&prompt)
25+
XCTAssertEqual(scopes, [])
26+
XCTAssertEqual(prompt, prompt)
27+
}
28+
29+
func test_parse_multiple_scopes() async throws {
30+
var prompt = "@web+file+selection hello"
31+
let scopes = parse(&prompt)
32+
XCTAssertEqual(scopes, ["web", "file", "selection"])
33+
XCTAssertEqual(prompt, "hello")
34+
}
35+
}
36+
37+
38+
39+

‎TestPlan.xctestplan‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,13 @@
7777
"identifier" : "OpenAIServiceTests",
7878
"name" : "OpenAIServiceTests"
7979
}
80+
},
81+
{
82+
"target" : {
83+
"containerPath" : "container:Core",
84+
"identifier" : "ChatServiceTests",
85+
"name" : "ChatServiceTests"
86+
}
8087
}
8188
],
8289
"version" : 1

0 commit comments

Comments
 (0)