Skip to content

Commit 3ef796a

Browse files
intitniintitni
committed
Merge branch 'feature/response-parsing-for-various-cases' into develop
2 parents 2793267 + 443d6b0 commit 3ef796a

10 files changed

Lines changed: 468 additions & 148 deletions

‎Core/Sources/CodeCompletionService/OllamaService.swift‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,10 @@ public actor OllamaService {
1111
let stopWords: [String]
1212
let keepAlive: String
1313
let format: ResponseFormat
14-
14+
1515
public enum ResponseFormat: String {
1616
case none = ""
17-
case json = "json"
17+
case json
1818
}
1919

2020
public enum Endpoint {
@@ -192,7 +192,13 @@ extension OllamaService {
192192
throw Error.otherError(text)
193193
}
194194

195-
return ResponseStream(result: result)
195+
return ResponseStream(result: result) {
196+
let chunk = try JSONDecoder().decode(
197+
ChatCompletionResponseChunk.self,
198+
from: $0.data(using: .utf8) ?? Data()
199+
)
200+
return .init(chunk: chunk, done: chunk.done)
201+
}
196202
}
197203
}
198204

@@ -249,7 +255,13 @@ extension OllamaService {
249255
throw Error.otherError(text)
250256
}
251257

252-
return ResponseStream(result: result)
258+
return ResponseStream(result: result) {
259+
let chunk = try JSONDecoder().decode(
260+
ChatCompletionResponseChunk.self,
261+
from: $0.data(using: .utf8) ?? Data()
262+
)
263+
return .init(chunk: chunk, done: chunk.done)
264+
}
253265
}
254266

255267
func countToken(_ message: Message) -> Int {

‎Core/Sources/CodeCompletionService/ResponseStream.swift‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,33 @@
11
import Foundation
22

3-
struct ResponseStream<Chunk: Decodable>: AsyncSequence {
3+
struct ResponseStream<Chunk>: AsyncSequence {
44
func makeAsyncIterator() -> Stream.AsyncIterator {
55
stream.makeAsyncIterator()
66
}
77

88
typealias Stream = AsyncThrowingStream<Chunk, Error>
99
typealias AsyncIterator = Stream.AsyncIterator
1010
typealias Element = Chunk
11+
12+
struct LineContent {
13+
let chunk: Chunk?
14+
let done: Bool
15+
}
1116

1217
let stream: Stream
1318

14-
init(result: URLSession.AsyncBytes, lineExtractor: @escaping (String) -> String? = { $0 }) {
19+
init(result: URLSession.AsyncBytes, lineExtractor: @escaping (String) throws -> LineContent) {
1520
stream = AsyncThrowingStream<Chunk, Error> { continuation in
1621
let task = Task {
1722
do {
1823
for try await line in result.lines {
1924
if Task.isCancelled { break }
20-
guard let content = lineExtractor(line)?.data(using: .utf8)
21-
else { continue }
22-
let chunk = try JSONDecoder().decode(Chunk.self, from: content)
23-
continuation.yield(chunk)
25+
let content = try lineExtractor(line)
26+
if let chunk = content.chunk {
27+
continuation.yield(chunk)
28+
}
29+
30+
if content.done { break }
2431
}
2532
continuation.finish()
2633
} catch {

‎Core/Sources/SuggestionService/RawSuggestionPostProcessing/DefaultRawSuggestionPostProcessingStrategy.swift‎

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
import Foundation
2+
import Parsing
3+
4+
protocol RawSuggestionPostProcessingStrategy {
5+
func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String
6+
}
7+
8+
extension RawSuggestionPostProcessingStrategy {
9+
func removeTrailingNewlinesAndWhitespace(from string: String) -> String {
10+
var text = string[...]
11+
while let last = text.last, last.isNewline || last.isWhitespace {
12+
text = text.dropLast(1)
13+
}
14+
return String(text)
15+
}
16+
}
17+
18+
struct DefaultRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy {
19+
let openingCodeTag: String
20+
let closingCodeTag: String
21+
22+
func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String {
23+
var suggestion = extractSuggestion(from: rawSuggestion)
24+
removePrefix(from: &suggestion, infillPrefix: infillPrefix)
25+
removeSuffix(from: &suggestion, suffix: suffix)
26+
return removeTrailingNewlinesAndWhitespace(from: infillPrefix + suggestion)
27+
}
28+
29+
func extractSuggestion(from response: String) -> String {
30+
let escapedMarkdownCodeBlock = removeLeadingAndTrailingMarkdownCodeBlockMark(from: response)
31+
let escapedTags = extractEnclosingSuggestion(
32+
from: escapedMarkdownCodeBlock,
33+
openingTag: openingCodeTag,
34+
closingTag: closingCodeTag
35+
)
36+
37+
return escapedTags
38+
}
39+
40+
func removePrefix(from suggestion: inout String, infillPrefix: String) {
41+
if suggestion.hasPrefix(infillPrefix) {
42+
suggestion.removeFirst(infillPrefix.count)
43+
}
44+
}
45+
46+
/// Window-mapping the lines in suggestion and the suffix to remove the common suffix.
47+
func removeSuffix(from suggestion: inout String, suffix: [String]) {
48+
let suggestionLines = suggestion.breakLines(appendLineBreakToLastLine: true)
49+
if let last = suggestionLines.last, let lastIndex = suffix.firstIndex(of: last) {
50+
var i = lastIndex - 1
51+
var j = suggestionLines.endIndex - 2
52+
while i >= 0, j >= 0, suffix[i] == suggestionLines[j] {
53+
i -= 1
54+
j -= 1
55+
}
56+
if i < 0 {
57+
let endIndex = max(j, 0)
58+
suggestion = suggestionLines[...endIndex].joined()
59+
}
60+
}
61+
}
62+
63+
/// Extract suggestions that is enclosed in tags.
64+
fileprivate func extractEnclosingSuggestion(
65+
from response: String,
66+
openingTag: String,
67+
closingTag: String
68+
) -> String {
69+
let case_openingTagAtTheStart_parseEverythingInsideTheTag = Parse(input: Substring.self) {
70+
openingTag
71+
72+
OneOf { // parse until tags or the end
73+
Parse {
74+
OneOf {
75+
PrefixUpTo(openingTag)
76+
PrefixUpTo(closingTag)
77+
}
78+
Skip {
79+
Rest()
80+
}
81+
}
82+
83+
Rest()
84+
}
85+
}
86+
87+
let case_noTagAtTheStart_parseEverythingBeforeTheTag = Parse(input: Substring.self) {
88+
OneOf {
89+
PrefixUpTo(openingTag)
90+
PrefixUpTo(closingTag)
91+
}
92+
93+
Skip {
94+
Rest()
95+
}
96+
}
97+
98+
let parser = Parse(input: Substring.self) {
99+
OneOf {
100+
case_openingTagAtTheStart_parseEverythingInsideTheTag
101+
case_noTagAtTheStart_parseEverythingBeforeTheTag
102+
Rest()
103+
}
104+
}
105+
106+
var text = response[...]
107+
do {
108+
let suggestion = try parser.parse(&text)
109+
return String(suggestion)
110+
} catch {
111+
return response
112+
}
113+
}
114+
115+
/// If the response starts with markdown code block, we should remove it.
116+
fileprivate func removeLeadingAndTrailingMarkdownCodeBlockMark(from response: String)
117+
-> String
118+
{
119+
let leadingMarkdownCodeBlockMarkParser = Parse(input: Substring.self) {
120+
Skip {
121+
Many {
122+
OneOf {
123+
" "
124+
"\n"
125+
}
126+
}
127+
"```"
128+
}
129+
}
130+
131+
let messagePrefixingMarkdownCodeBlockMarkParser = Parse(input: Substring.self) {
132+
Skip {
133+
PrefixThrough(":")
134+
"\n```"
135+
}
136+
}
137+
138+
let removePrefixMarkdownCodeBlockMark = Parse(input: Substring.self) {
139+
Skip {
140+
OneOf {
141+
leadingMarkdownCodeBlockMarkParser
142+
messagePrefixingMarkdownCodeBlockMarkParser
143+
}
144+
PrefixThrough("\n")
145+
}
146+
OneOf {
147+
Parse {
148+
PrefixUpTo("```")
149+
Skip { Rest() }
150+
}
151+
Rest()
152+
}
153+
}
154+
155+
do {
156+
var response = response[...]
157+
let suggestion = try removePrefixMarkdownCodeBlockMark.parse(&response)
158+
return String(suggestion)
159+
} catch {
160+
return response
161+
}
162+
}
163+
}
164+

‎Core/Sources/SuggestionService/RawSuggestionPostProcessing/NoOpRawSuggestionPostProcessingStrategy.swift‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import Foundation
2+
3+
struct NoOpRawSuggestionPostProcessingStrategy: RawSuggestionPostProcessingStrategy {
4+
func postProcess(rawSuggestion: String, infillPrefix: String, suffix: [String]) -> String {
5+
removeTrailingNewlinesAndWhitespace(from: infillPrefix + rawSuggestion)
6+
}
7+
}
8+

‎Core/Sources/SuggestionService/RequestStrategies/ContinueRequestStrategy.swift‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,13 @@ struct ContinueRequestStrategy: RequestStrategy {
2121
suffix: suffix
2222
)
2323
}
24+
25+
func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
26+
DefaultRawSuggestionPostProcessingStrategy(
27+
openingCodeTag: Tag.openingCode,
28+
closingCodeTag: Tag.closingCode
29+
)
30+
}
2431

2532
enum Tag {
2633
public static let openingCode = "<Code3721>"
@@ -158,21 +165,5 @@ struct ContinueRequestStrategy: RequestStrategy {
158165
)
159166
}
160167
}
161-
162-
func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String {
163-
let suggestion = extractEnclosingSuggestion(
164-
from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion),
165-
openingTag: Tag.openingCode,
166-
closingTag: Tag.closingCode
167-
)
168-
169-
if suggestion.hasPrefix(suggestionPrefix) {
170-
var processed = suggestion
171-
processed.removeFirst(suggestionPrefix.count)
172-
return processed
173-
}
174-
175-
return suggestionPrefix + suggestion
176-
}
177168
}
178169

‎Core/Sources/SuggestionService/RequestStrategies/DefaultRequestStrategy.swift‎

Lines changed: 7 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,13 @@ struct DefaultRequestStrategy: RequestStrategy {
2222
suffix: suffix
2323
)
2424
}
25+
26+
func createRawSuggestionPostProcessor() -> DefaultRawSuggestionPostProcessingStrategy {
27+
DefaultRawSuggestionPostProcessingStrategy(
28+
openingCodeTag: Tag.openingCode,
29+
closingCodeTag: Tag.closingCode
30+
)
31+
}
2532

2633
enum Tag {
2734
public static let openingCode = "<Code3721>"
@@ -142,21 +149,5 @@ struct DefaultRequestStrategy: RequestStrategy {
142149
)
143150
}
144151
}
145-
146-
func postProcessRawSuggestion(suggestionPrefix: String, suggestion: String) -> String {
147-
let suggestion = extractEnclosingSuggestion(
148-
from: removeLeadingAndTrailingMarkdownCodeBlockMark(from: suggestion),
149-
openingTag: Tag.openingCode,
150-
closingTag: Tag.closingCode
151-
)
152-
153-
if suggestion.hasPrefix(suggestionPrefix) {
154-
var processed = suggestion
155-
processed.removeFirst(suggestionPrefix.count)
156-
return processed
157-
}
158-
159-
return suggestionPrefix + suggestion
160-
}
161152
}
162153

‎Core/Sources/SuggestionService/RequestStrategies/NaiveRequestStrategy.swift‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ struct NaiveRequestStrategy: RequestStrategy {
2020
suffix: suffix
2121
)
2222
}
23+
24+
func createRawSuggestionPostProcessor() -> some RawSuggestionPostProcessingStrategy {
25+
NoOpRawSuggestionPostProcessingStrategy()
26+
}
2327

2428
struct Request: PromptStrategy {
2529
let systemPrompt: String = ""

0 commit comments

Comments
 (0)