@@ -3,9 +3,9 @@ import Fundamental
33import Storage
44
55protocol CodeCompletionServiceType {
6- func getCompletion (
7- _ request : PromptStrategy
8- ) async throws -> String
6+ associatedtype CompletionSequence : AsyncSequence where CompletionSequence . Element == String
7+
8+ func getCompletion ( _ request : PromptStrategy ) async throws -> CompletionSequence
99}
1010
1111extension CodeCompletionServiceType {
@@ -16,7 +16,12 @@ extension CodeCompletionServiceType {
1616 try await withThrowingTaskGroup ( of: String . self) { group in
1717 for _ in 0 ..< max ( 1 , count) {
1818 _ = group. addTaskUnlessCancelled {
19- try await getCompletion ( request)
19+ var result = " "
20+ let stream = try await getCompletion ( request)
21+ for try await response in stream {
22+ result. append ( response)
23+ }
24+ return result
2025 }
2126 }
2227
@@ -110,6 +115,18 @@ public struct CodeCompletionService {
110115 let result = try await service. getCompletions ( prompt, count: count)
111116 try Task . checkCancellation ( )
112117 return result
118+ case . ollama:
119+ let service = OllamaService (
120+ url: model. endpoint,
121+ endpoint: . chatCompletion,
122+ modelName: model. info. modelName,
123+ stopWords: prompt. stopWords,
124+ keepAlive: model. info. ollamaInfo. keepAlive,
125+ format: . none
126+ )
127+ let result = try await service. getCompletions ( prompt, count: count)
128+ try Task . checkCancellation ( )
129+ return result
113130 case . unknown:
114131 throw Error . unknownFormat
115132 }
@@ -145,6 +162,18 @@ public struct CodeCompletionService {
145162 let result = try await service. getCompletions ( prompt, count: count)
146163 try Task . checkCancellation ( )
147164 return result
165+ case . ollama:
166+ let service = OllamaService (
167+ url: model. endpoint,
168+ endpoint: . completion,
169+ modelName: model. info. modelName,
170+ stopWords: prompt. stopWords,
171+ keepAlive: model. info. ollamaInfo. keepAlive,
172+ format: . none
173+ )
174+ let result = try await service. getCompletions ( prompt, count: count)
175+ try Task . checkCancellation ( )
176+ return result
148177 case . unknown:
149178 throw Error . unknownFormat
150179 }
0 commit comments