Skip to content

Commit e6dee6b

Browse files
intitniintitni
committed
Add header value parser
1 parent d8f95d2 commit e6dee6b

6 files changed

Lines changed: 140 additions & 20 deletions

‎Tool/Sources/OpenAIService/APIs/ChatCompletionsAPIDefinition.swift‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,23 @@ protocol ChatCompletionsStreamAPI {
172172
func callAsFunction() async throws -> AsyncThrowingStream<ChatCompletionsStreamDataChunk, Error>
173173
}
174174

175+
extension ChatCompletionsStreamAPI {
176+
static func setupExtraHeaderFields(
177+
_ request: inout URLRequest,
178+
model: ChatModel,
179+
apiKey: String
180+
) async {
181+
let parser = HeaderValueParser()
182+
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
183+
let value = await parser.parse(
184+
field.value,
185+
context: .init(modelName: model.info.modelName, apiKey: apiKey)
186+
)
187+
request.setValue(value, forHTTPHeaderField: field.key)
188+
}
189+
}
190+
}
191+
175192
extension AsyncSequence {
176193
func toStream() -> AsyncThrowingStream<Element, Error> {
177194
AsyncThrowingStream { continuation in

‎Tool/Sources/OpenAIService/APIs/EmbeddingAPIDefinitions.swift‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,23 @@ protocol EmbeddingAPI {
99
func embed(tokens: [[Int]]) async throws -> EmbeddingResponse
1010
}
1111

12+
extension EmbeddingAPI {
13+
static func setupExtraHeaderFields(
14+
_ request: inout URLRequest,
15+
model: EmbeddingModel,
16+
apiKey: String
17+
) async {
18+
let parser = HeaderValueParser()
19+
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
20+
let value = await parser.parse(
21+
field.value,
22+
context: .init(modelName: model.info.modelName, apiKey: apiKey)
23+
)
24+
request.setValue(value, forHTTPHeaderField: field.key)
25+
}
26+
}
27+
}
28+
1229
public struct EmbeddingResponse: Decodable {
1330
public struct Object: Decodable {
1431
public var embedding: [Float]

‎Tool/Sources/OpenAIService/APIs/OlamaChatCompletionsService.swift‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,7 @@ extension OllamaChatCompletionsService: ChatCompletionsAPI {
6464
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
6565
}
6666

67-
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
68-
request.setValue(field.value, forHTTPHeaderField: field.key)
69-
}
67+
await Self.setupExtraHeaderFields(&request, model: model, apiKey: apiKey)
7068

7169
let (result, response) = try await URLSession.shared.data(for: request)
7270

@@ -149,9 +147,7 @@ extension OllamaChatCompletionsService: ChatCompletionsStreamAPI {
149147
request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
150148
}
151149

152-
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
153-
request.setValue(field.value, forHTTPHeaderField: field.key)
154-
}
150+
await Self.setupExtraHeaderFields(&request, model: model, apiKey: apiKey)
155151

156152
let (result, response) = try await URLSession.shared.bytes(for: request)
157153

‎Tool/Sources/OpenAIService/APIs/OpenAIChatCompletionsService.swift‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,7 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
327327

328328
Self.setupAppInformation(&request)
329329
Self.setupAPIKey(&request, model: model, apiKey: apiKey)
330-
Self.setupExtraHeaderFields(&request, model: model)
330+
await Self.setupExtraHeaderFields(&request, model: model, apiKey: apiKey)
331331
requestModifier?(&request)
332332

333333
let (result, response) = try await URLSession.shared.bytes(for: request)
@@ -473,12 +473,6 @@ actor OpenAIChatCompletionsService: ChatCompletionsStreamAPI, ChatCompletionsAPI
473473
}
474474
}
475475
}
476-
477-
static func setupExtraHeaderFields(_ request: inout URLRequest, model: ChatModel) {
478-
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
479-
request.setValue(field.value, forHTTPHeaderField: field.key)
480-
}
481-
}
482476
}
483477

484478
extension OpenAIChatCompletionsService.ResponseBody {

‎Tool/Sources/OpenAIService/APIs/OpenAIEmbeddingService.swift‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
4141

4242
Self.setupAppInformation(&request)
4343
Self.setupAPIKey(&request, model: model, apiKey: apiKey)
44+
await Self.setupExtraHeaderFields(&request, model: model, apiKey: apiKey)
4445

4546
let (result, response) = try await URLSession.shared.data(for: request)
4647
guard let response = response as? HTTPURLResponse else {
@@ -80,7 +81,7 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
8081

8182
Self.setupAppInformation(&request)
8283
Self.setupAPIKey(&request, model: model, apiKey: apiKey)
83-
Self.setupExtraHeaderFields(&request, model: model)
84+
await Self.setupExtraHeaderFields(&request, model: model, apiKey: apiKey)
8485

8586
let (result, response) = try await URLSession.shared.data(for: request)
8687
guard let response = response as? HTTPURLResponse else {
@@ -140,11 +141,5 @@ struct OpenAIEmbeddingService: EmbeddingAPI {
140141
}
141142
}
142143
}
143-
144-
static func setupExtraHeaderFields(_ request: inout URLRequest, model: EmbeddingModel) {
145-
for field in model.info.customHeaderInfo.headers where !field.key.isEmpty {
146-
request.setValue(field.value, forHTTPHeaderField: field.key)
147-
}
148-
}
149144
}
150145

‎Tool/Sources/OpenAIService/HeaderValueParser.swift‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import Foundation
2+
import GitHubCopilotService
3+
import Logger
4+
import Terminal
5+
6+
public struct HeaderValueParser {
7+
public enum Placeholder: String {
8+
case gitHubCopilotOBearerToken = "github_copilot_bearer_token"
9+
case apiKey = "api_key"
10+
case modelName = "model_name"
11+
}
12+
13+
public struct Context {
14+
public var modelName: String
15+
public var apiKey: String
16+
public var gitHubCopilotToken: () async -> GitHubCopilotExtension.Token?
17+
public var shellEnvironmentVariable: (_ key: String) async -> String?
18+
19+
public init(
20+
modelName: String,
21+
apiKey: String,
22+
gitHubCopilotToken: (() async -> GitHubCopilotExtension.Token?)? = nil,
23+
shellEnvironmentVariable: ((_: String) async -> String?)? = nil
24+
) {
25+
self.modelName = modelName
26+
self.apiKey = apiKey
27+
self.gitHubCopilotToken = gitHubCopilotToken ?? {
28+
try? await GitHubCopilotExtension.fetchToken()
29+
}
30+
self.shellEnvironmentVariable = shellEnvironmentVariable ?? { p in
31+
let shell = ProcessInfo.processInfo.environment["SHELL"] ?? "/bin/bash"
32+
let terminal = Terminal()
33+
return try? await terminal.runCommand(
34+
shell,
35+
arguments: ["-i", "-l", "-c", "echo $\(p)"],
36+
environment: [:]
37+
)
38+
}
39+
}
40+
}
41+
42+
public init() {}
43+
44+
/// Replace `{{PlaceHolder}}` with exact values.
45+
public func parse(_ value: String, context: Context) async -> String {
46+
var parsedValue = value
47+
let placeholderRanges = findPlaceholderRanges(in: parsedValue)
48+
49+
for (range, placeholderText) in placeholderRanges.reversed() {
50+
let cleanPlaceholder = placeholderText
51+
.trimmingCharacters(in: CharacterSet(charactersIn: "{}"))
52+
53+
var replacement: String?
54+
if let knownPlaceholder = Placeholder(rawValue: cleanPlaceholder) {
55+
async let token = context.gitHubCopilotToken()
56+
switch knownPlaceholder {
57+
case .gitHubCopilotOBearerToken:
58+
replacement = await token?.token
59+
case .apiKey:
60+
replacement = context.apiKey
61+
case .modelName:
62+
replacement = context.modelName
63+
}
64+
} else {
65+
replacement = await context.shellEnvironmentVariable(cleanPlaceholder)
66+
}
67+
68+
if let replacement {
69+
parsedValue.replaceSubrange(range, with: replacement)
70+
} else {
71+
parsedValue.replaceSubrange(range, with: "none")
72+
}
73+
}
74+
75+
return parsedValue
76+
}
77+
78+
private func findPlaceholderRanges(in string: String) -> [(Range<String.Index>, String)] {
79+
var ranges: [(Range<String.Index>, String)] = []
80+
let pattern = #"\{\{[^}]+\}\}"#
81+
82+
do {
83+
let regex = try NSRegularExpression(pattern: pattern)
84+
let matches = regex.matches(
85+
in: string,
86+
range: NSRange(string.startIndex..., in: string)
87+
)
88+
89+
for match in matches {
90+
if let range = Range(match.range, in: string) {
91+
ranges.append((range, String(string[range])))
92+
}
93+
}
94+
} catch {
95+
Logger.service.error("Failed to find placeholders in string: \(string)")
96+
}
97+
98+
return ranges
99+
}
100+
}
101+

0 commit comments

Comments
 (0)