Skip to content

Commit 0ccb0c1

Browse files
intitniintitni
committed
Add SerpAPI
1 parent ca0bd7a commit 0ccb0c1

7 files changed

Lines changed: 217 additions & 38 deletions

File tree

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

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
import BingSearchService
21
import ChatBasic
32
import Foundation
43
import OpenAIService
54
import Preferences
5+
import WebSearchService
66

77
struct SearchFunction: ChatGPTFunction {
88
static let dateFormatter = {
@@ -17,13 +17,13 @@ struct SearchFunction: ChatGPTFunction {
1717
}
1818

1919
struct Result: ChatGPTFunctionResult {
20-
var result: BingSearchResult
20+
var result: WebSearchResult
2121

2222
var botReadableContent: String {
23-
result.webPages.value.enumerated().map {
23+
result.webPages.enumerated().map {
2424
let (index, page) = $0
2525
return """
26-
\(index + 1). \(page.name) \(page.url)
26+
\(index + 1). \(page.title) \(page.urlString)
2727
\(page.snippet)
2828
"""
2929
}.joined(separator: "\n")
@@ -72,22 +72,15 @@ struct SearchFunction: ChatGPTFunction {
7272
await reportProgress("Searching \(arguments.query)")
7373

7474
do {
75-
let bingSearch = BingSearchService(
76-
subscriptionKey: UserDefaults.shared.value(for: \.bingSearchSubscriptionKey),
77-
searchURL: UserDefaults.shared.value(for: \.bingSearchEndpoint)
78-
)
75+
let search = WebSearchService(provider: .userPreferred)
7976

80-
let result = try await bingSearch.search(
81-
query: arguments.query,
82-
numberOfResult: maxTokens > 5000 ? 5 : 3,
83-
freshness: arguments.freshness
84-
)
77+
let result = try await search.search(query: arguments.query)
8578

8679
await reportProgress("""
8780
Finish searching \(arguments.query)
8881
\(
89-
result.webPages.value
90-
.map { "- [\($0.name)](\($0.url))" }
82+
result.webPages
83+
.map { "- [\($0.title)](\($0.urlString))" }
9184
.joined(separator: "\n")
9285
)
9386
""")

Core/Sources/HostApp/AccountSettings/BingSearchView.swift renamed to Core/Sources/HostApp/AccountSettings/WebSearchView.swift

File renamed without changes.

‎Tool/Package.swift‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ let package = Package(
1010
.library(name: "XPCShared", targets: ["XPCShared"]),
1111
.library(name: "Terminal", targets: ["Terminal"]),
1212
.library(name: "LangChain", targets: ["LangChain"]),
13-
.library(name: "ExternalServices", targets: ["BingSearchService"]),
13+
.library(name: "ExternalServices", targets: ["WebSearchService"]),
1414
.library(name: "Preferences", targets: ["Preferences", "Configs"]),
1515
.library(name: "Logger", targets: ["Logger"]),
1616
.library(name: "OpenAIService", targets: ["OpenAIService"]),
@@ -52,7 +52,7 @@ let package = Package(
5252
.library(name: "CommandHandler", targets: ["CommandHandler"]),
5353
.library(name: "CodeDiff", targets: ["CodeDiff"]),
5454
.library(name: "BuiltinExtension", targets: ["BuiltinExtension"]),
55-
.library(name: "BingSearchService", targets: ["BingSearchService"]),
55+
.library(name: "WebSearchService", targets: ["WebSearchService"]),
5656
.library(
5757
name: "CustomCommandTemplateProcessor",
5858
targets: ["CustomCommandTemplateProcessor"]
@@ -382,7 +382,7 @@ let package = Package(
382382
]
383383
),
384384

385-
.target(name: "BingSearchService"),
385+
.target(name: "WebSearchService", dependencies: ["Preferences"]),
386386

387387
.target(name: "SuggestionProvider", dependencies: [
388388
"SuggestionBasic",

‎Tool/Sources/Preferences/Keys.swift‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -582,14 +582,43 @@ public extension UserDefaultPreferenceKeys {
582582
}
583583
}
584584

585-
// MARK: - Bing Search
585+
// MARK: - Search
586586

587587
public extension UserDefaultPreferenceKeys {
588-
var bingSearchSubscriptionKey: PreferenceKey<String> {
588+
enum SearchProvider: String, Codable, CaseIterable {
589+
case serpAPI
590+
case headlessBrowser
591+
}
592+
593+
enum SerpAPIEngine: String, Codable, CaseIterable {
594+
case google
595+
case baidu
596+
case bing
597+
case duckDuckGo = "duckduckgo"
598+
}
599+
600+
enum HeadlessBrowserEngine: String, Codable, CaseIterable {
601+
case google
602+
case baidu
603+
}
604+
605+
var searchProvider: PreferenceKey<SearchProvider> {
606+
.init(defaultValue: .headlessBrowser, key: "SearchProvider")
607+
}
608+
609+
var serpAPIEngine: PreferenceKey<SerpAPIEngine> {
610+
.init(defaultValue: .google, key: "SerpAPIEngine")
611+
}
612+
613+
var headlessBrowserEngine: PreferenceKey<HeadlessBrowserEngine> {
614+
.init(defaultValue: .google, key: "HeadlessBrowserEngine")
615+
}
616+
617+
var bingSearchSubscriptionKey: DeprecatedPreferenceKey<String> {
589618
.init(defaultValue: "", key: "BingSearchSubscriptionKey")
590619
}
591620

592-
var bingSearchEndpoint: PreferenceKey<String> {
621+
var bingSearchEndpoint: DeprecatedPreferenceKey<String> {
593622
.init(
594623
defaultValue: "https://api.bing.microsoft.com/v7.0/search/",
595624
key: "BingSearchEndpoint"

Tool/Sources/BingSearchService/BingSearchService.swift renamed to Tool/Sources/WebSearchService/SearchServices/BingSearchService.swift

Lines changed: 28 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,19 @@
11
import Foundation
22

3-
public struct BingSearchResult: Codable {
4-
public var webPages: WebPages
3+
struct BingSearchResult: Codable {
4+
var webPages: WebPages
55

6-
public struct WebPages: Codable {
7-
public var webSearchUrl: String
8-
public var totalEstimatedMatches: Int
9-
public var value: [WebPageValue]
6+
struct WebPages: Codable {
7+
var webSearchUrl: String
8+
var totalEstimatedMatches: Int
9+
var value: [WebPageValue]
1010

11-
public struct WebPageValue: Codable {
12-
public var id: String
13-
public var name: String
14-
public var url: String
15-
public var displayUrl: String
16-
public var snippet: String
11+
struct WebPageValue: Codable {
12+
var id: String
13+
var name: String
14+
var url: String
15+
var displayUrl: String
16+
var snippet: String
1717
}
1818
}
1919
}
@@ -42,16 +42,27 @@ enum BingSearchError: Error, LocalizedError {
4242
}
4343
}
4444

45-
public struct BingSearchService {
46-
public var subscriptionKey: String
47-
public var searchURL: String
45+
struct BingSearchService: SearchService {
46+
var subscriptionKey: String
47+
var searchURL: String
4848

49-
public init(subscriptionKey: String, searchURL: String) {
49+
init(subscriptionKey: String, searchURL: String) {
5050
self.subscriptionKey = subscriptionKey
5151
self.searchURL = searchURL
5252
}
5353

54-
public func search(
54+
func search(query: String) async throws -> WebSearchResult {
55+
let result = try await search(query: query, numberOfResult: 10)
56+
return WebSearchResult(webPages: result.webPages.value.map {
57+
WebSearchResult.WebPage(
58+
urlString: $0.url,
59+
title: $0.name,
60+
snippet: $0.snippet
61+
)
62+
})
63+
}
64+
65+
func search(
5566
query: String,
5667
numberOfResult: Int,
5768
freshness: String? = nil

‎Tool/Sources/WebSearchService/SearchServices/SerpAPISearchService.swift‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import Foundation
2+
3+
struct SerpAPIResponse: Codable {
4+
var searchMetadata: SearchMetadata
5+
var organicResults: [OrganicResult]
6+
7+
struct SearchMetadata: Codable {
8+
var id: String
9+
var status: String
10+
var jsonEndpoint: String
11+
var createdAt: String
12+
var processedAt: String
13+
var totalTimeTaken: Double
14+
}
15+
16+
struct OrganicResult: Codable {
17+
var position: Int
18+
var title: String
19+
var link: String
20+
var snippet: String
21+
22+
func toWebSearchResult() -> WebSearchResult.WebPage {
23+
return WebSearchResult.WebPage(urlString: link, title: title, snippet: snippet)
24+
}
25+
}
26+
27+
func toWebSearchResult() -> WebSearchResult {
28+
return WebSearchResult(webPages: organicResults.map { $0.toWebSearchResult() })
29+
}
30+
}
31+
32+
struct SerpAPISearchService: SearchService {
33+
let engine: WebSearchProvider.SerpAPIEngine
34+
let endpoint: URL = .init(string: "https://serpapi.com/search.json")!
35+
let apiKey: String
36+
37+
init(engine: WebSearchProvider.SerpAPIEngine, apiKey: String) {
38+
self.engine = engine
39+
self.apiKey = apiKey
40+
}
41+
42+
func search(query: String) async throws -> WebSearchResult {
43+
var request = URLRequest(url: endpoint)
44+
request.httpMethod = "GET"
45+
var urlComponents = URLComponents(url: endpoint, resolvingAgainstBaseURL: true)!
46+
urlComponents.queryItems = [
47+
URLQueryItem(name: "q", value: query),
48+
URLQueryItem(name: "engine", value: engine.rawValue),
49+
URLQueryItem(name: "api_key", value: apiKey)
50+
]
51+
52+
guard let url = urlComponents.url else {
53+
throw URLError(.badURL)
54+
}
55+
56+
request = URLRequest(url: url)
57+
request.httpMethod = "GET"
58+
59+
let (data, response) = try await URLSession.shared.data(for: request)
60+
61+
guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else {
62+
throw URLError(.badServerResponse)
63+
}
64+
65+
// Parse the response into WebSearchResult
66+
let decoder = JSONDecoder()
67+
68+
do {
69+
let searchResponse = try decoder.decode(SerpAPIResponse.self, from: data)
70+
return searchResponse.toWebSearchResult()
71+
} catch {
72+
throw error
73+
}
74+
}
75+
}
76+

‎Tool/Sources/WebSearchService/WebSearchService.swift‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import Foundation
2+
import Preferences
3+
4+
public enum WebSearchProvider {
5+
public enum SerpAPIEngine: String {
6+
case google
7+
case baidu
8+
case bing
9+
case duckDuckGo = "duckduckgo"
10+
}
11+
12+
public enum HeadlessBrowserEngine: String {
13+
case google
14+
case baidu
15+
}
16+
17+
case serpAPI(SerpAPIEngine, apiKey: String)
18+
case headlessBrowser(HeadlessBrowserEngine)
19+
20+
public static var userPreferred: WebSearchProvider {
21+
switch UserDefaults.shared.value(for: \.searchProvider) {
22+
case .headlessBrowser:
23+
return .headlessBrowser(.init(
24+
rawValue: UserDefaults.shared.value(for: \.headlessBrowserEngine)
25+
.rawValue
26+
) ?? .google)
27+
case .serpAPI:
28+
return .serpAPI(.init(
29+
rawValue: UserDefaults.shared.value(for: \.headlessBrowserEngine)
30+
.rawValue
31+
) ?? .google, apiKey: "")
32+
}
33+
}
34+
}
35+
36+
public struct WebSearchResult {
37+
public struct WebPage {
38+
public var urlString: String
39+
public var title: String
40+
public var snippet: String
41+
}
42+
43+
public var webPages: [WebPage]
44+
}
45+
46+
public protocol SearchService {
47+
func search(query: String) async throws -> WebSearchResult
48+
}
49+
50+
public struct WebSearchService {
51+
let service: SearchService
52+
53+
init(service: SearchService) {
54+
self.service = service
55+
}
56+
57+
public init(provider: WebSearchProvider) {
58+
switch provider {
59+
case let .serpAPI(engine, apiKey):
60+
service = SerpAPISearchService(engine: engine, apiKey: apiKey)
61+
case let .headlessBrowser(engine):
62+
fatalError()
63+
}
64+
}
65+
66+
public func search(query: String) async throws -> WebSearchResult {
67+
return try await service.search(query: query)
68+
}
69+
}
70+

0 commit comments

Comments
 (0)