Skip to content

Commit fb6be98

Browse files
intitniintitni
committed
Add CodeiumService
1 parent 1ed614d commit fb6be98

6 files changed

Lines changed: 742 additions & 0 deletions

File tree

‎Core/Package.swift‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,7 @@ let package = Package(
117117
),
118118
.target(name: "SuggestionService", dependencies: [
119119
"GitHubCopilotService",
120+
"CodeiumService",
120121
]),
121122

122123
// MARK: - Prompt To Code
@@ -205,6 +206,13 @@ let package = Package(
205206
name: "OpenAIServiceTests",
206207
dependencies: ["OpenAIService"]
207208
),
209+
210+
// MARK: - Codeium
211+
212+
.target(
213+
name: "CodeiumService",
214+
dependencies: ["LanguageClient", "SuggestionModel", "Preferences"]
215+
),
208216
]
209217
)
210218

‎Core/Sources/CodeiumService/CodeiumLanguageServer.swift‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
import Foundation
2+
import JSONRPC
3+
import LanguageClient
4+
import LanguageServerProtocol
5+
6+
protocol CodeiumLSP {
7+
func sendRequest<E: CodeiumRequestType>(_ endpoint: E) async throws -> E.Response
8+
}
9+
10+
class CodeiumLanguageServer: CodeiumLSP {
11+
let languageServerExecutableURL: URL
12+
let managerDirectoryURL: URL
13+
let supportURL: URL
14+
let process: Process
15+
let transport: StdioDataTransport
16+
var terminationHandler: (() -> Void)?
17+
var launchHandler: (() -> Void)?
18+
var port: String?
19+
20+
init(
21+
languageServerExecutableURL: URL,
22+
managerDirectoryURL: URL,
23+
supportURL: URL,
24+
terminationHandler: (() -> Void)? = nil,
25+
launchHandler: (() -> Void)? = nil
26+
) {
27+
self.languageServerExecutableURL = languageServerExecutableURL
28+
self.managerDirectoryURL = managerDirectoryURL
29+
self.supportURL = supportURL
30+
self.terminationHandler = terminationHandler
31+
self.launchHandler = launchHandler
32+
process = Process()
33+
transport = StdioDataTransport()
34+
35+
process.standardInput = transport.stdinPipe
36+
process.standardOutput = transport.stdoutPipe
37+
process.standardError = transport.stderrPipe
38+
39+
process.executableURL = languageServerExecutableURL
40+
41+
process.arguments = [
42+
"--api_server_host",
43+
"server.codeium.com",
44+
"--api_server_port",
45+
"443",
46+
"--manager_dir",
47+
managerDirectoryURL.path,
48+
]
49+
50+
process.currentDirectoryURL = supportURL
51+
52+
process.terminationHandler = { [weak self] task in
53+
self?.processTerminated(task)
54+
}
55+
}
56+
57+
func start() {
58+
guard !process.isRunning else { return }
59+
port = nil
60+
do {
61+
try process.run()
62+
63+
Task {
64+
func findPort() -> String? {
65+
// find a file in managerDirectoryURL whose name looks like a port, return the
66+
// name if found
67+
let fileManager = FileManager.default
68+
let enumerator = fileManager.enumerator(
69+
at: managerDirectoryURL,
70+
includingPropertiesForKeys: nil
71+
)
72+
while let fileURL = enumerator?.nextObject() as? URL {
73+
if fileURL.lastPathComponent.range(
74+
of: #"^\d+$"#,
75+
options: .regularExpression
76+
) != nil {
77+
return fileURL.lastPathComponent
78+
}
79+
}
80+
return nil
81+
}
82+
83+
try await Task.sleep(nanoseconds: 2_000_000)
84+
port = findPort()
85+
var waited = 0
86+
87+
while true {
88+
try await Task.sleep(nanoseconds: 1_000_000_000)
89+
waited += 1
90+
port = findPort()
91+
if port != nil { break }
92+
if waited >= 60 {
93+
process.terminate()
94+
}
95+
}
96+
}
97+
} catch {
98+
print("start: ", error)
99+
processTerminated(process)
100+
}
101+
}
102+
103+
deinit {
104+
process.terminationHandler = nil
105+
process.terminate()
106+
transport.close()
107+
}
108+
109+
private func processTerminated(_: Process) {
110+
transport.close()
111+
terminationHandler?()
112+
}
113+
114+
func handleFileEvent(_ event: FileEvent) {
115+
switch event.type {
116+
case .created:
117+
let fileURL = URL(fileURLWithPath: event.uri)
118+
var isDirectory: ObjCBool = false
119+
if FileManager.default.fileExists(atPath: fileURL.path, isDirectory: &isDirectory),
120+
!isDirectory.boolValue
121+
{
122+
let portName = fileURL.lastPathComponent
123+
print("set port", portName)
124+
if port == nil {
125+
port = portName
126+
}
127+
}
128+
default:
129+
break
130+
}
131+
}
132+
133+
func sendRequest<E>(_ request: E) async throws -> E.Response where E: CodeiumRequestType {
134+
guard let port else { throw CancellationError() }
135+
136+
let request = request.makeURLRequest(server: "http://127.0.0.1:\(port)")
137+
let (data, response) = try await URLSession.shared.data(for: request)
138+
if (response as? HTTPURLResponse)?.statusCode == 200 {
139+
do {
140+
let response = try JSONDecoder().decode(E.Response.self, from: data)
141+
return response
142+
} catch {
143+
dump(error)
144+
throw error
145+
}
146+
} else {
147+
do {
148+
let error = try JSONDecoder().decode(CodeiumResponseError.self, from: data)
149+
throw error
150+
} catch {
151+
print(String(data: data, encoding: .utf8))
152+
throw error
153+
}
154+
}
155+
}
156+
}
157+

‎Core/Sources/CodeiumService/CodeiumModels.swift‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
import Foundation
2+
import JSONRPC
3+
import LanguageServerProtocol
4+
import SuggestionModel
5+
6+
struct CodeiumCompletion: Codable {
7+
var completionId: String
8+
var text: String
9+
var prefix: String?
10+
var stop: String?
11+
// var score: Double
12+
// var tokens: [UInt64]
13+
// var decodedTokens: [String]
14+
// var probabilities: [Double]
15+
// var adjustedProbabilities: [Double]
16+
// var generatedLength: Int
17+
}
18+
19+
struct CodeiumCompletionItem: Codable {
20+
var completion: CodeiumCompletion
21+
var suffix: Suffix?
22+
var range: Range
23+
var source: CompletionSource
24+
var completionParts: [CompletionPart]
25+
}
26+
27+
struct Suffix: Codable {
28+
/// Text to insert after the cursor when accepting the completion.
29+
var text: String
30+
/// Cursor position delta (as signed offset) from the end of the inserted
31+
/// completion (including the suffix).
32+
var deltaCursorOffset: Int
33+
}
34+
35+
struct Range: Codable {
36+
var startOffset: String?
37+
var endOffset: String?
38+
var startPosition: DocumentPosition?
39+
var endPosition: DocumentPosition?
40+
}
41+
42+
enum CompletionSource: String, Codable {
43+
case unspecified = "COMPLETION_SOURCE_UNSPECIFIED"
44+
case typingAsSuggested = "COMPLETION_SOURCE_TYPING_AS_SUGGESTED"
45+
case cache = "COMPLETION_SOURCE_CACHE"
46+
case network = "COMPLETION_SOURCE_NETWORK"
47+
}
48+
49+
/// Represents a contiguous part of the completion text that is not
50+
/// already in the document.
51+
struct CompletionPart: Codable {
52+
enum CompletionPartType: String, Codable {
53+
case unspecified = "COMPLETION_PART_TYPE_UNSPECIFIED"
54+
/// Single-line completion parts that appear within an existing line of text.
55+
case inline = "COMPLETION_PART_TYPE_INLINE"
56+
/// Possibly multi-line completion parts that appear below an existing line of text.
57+
case block = "COMPLETION_PART_TYPE_BLOCK"
58+
/// Like COMPLETION_PART_TYPE_INLINE, but overwrites the existing text.
59+
case inline_mask = "COMPLETION_PART_TYPE_INLINE_MASK"
60+
}
61+
62+
var text: String
63+
/// Offset in the original document where the part starts. For block
64+
/// parts, this is always the end of the line before the block.
65+
var offset: String
66+
var type: CompletionPartType
67+
/// The section of the original line that came before this part. Only valid for
68+
/// COMPLETION_PART_TYPE_INLINE.
69+
var prefix: String?
70+
/// In the case of COMPLETION_PART_TYPE_BLOCK, represents the line it is below.
71+
var line: String?
72+
}
73+
74+
struct CodeiumDocument: Codable {
75+
var absolute_path: String
76+
// Path relative to the root of the workspace.
77+
var relative_path: String
78+
var text: String
79+
// Language ID provided by the editor.
80+
var editor_language: String
81+
// Language enum standardized across editors.
82+
var language: CodeiumSupportedLanguage
83+
// Measured in number of UTF-8 bytes.
84+
// var cursor_offset: UInt64?
85+
// May be present instead of cursor_offset.
86+
var cursor_position: DocumentCursorPosition?
87+
// \n or \r\n, if known.
88+
var line_ending: String = "\n"
89+
}
90+
91+
struct DocumentPosition: Codable {
92+
/// 0-indexed. Measured in UTF-8 bytes.
93+
var row: String?
94+
/// 0-indexed. Measured in UTF-8 bytes.
95+
var col: String?
96+
}
97+
98+
struct DocumentCursorPosition: Codable {
99+
/// 0-indexed. Measured in UTF-8 bytes.
100+
var row: Int
101+
/// 0-indexed. Measured in UTF-8 bytes.
102+
var col: Int
103+
}
104+
105+
struct CodeiumEditorOptions: Codable {
106+
var tab_size: Int
107+
var insert_spaces: Bool
108+
}
109+
110+
struct Metadata: Codable {
111+
var ide_name: String
112+
var ide_version: String
113+
var extension_name: String
114+
var extension_version: String
115+
var api_key: String
116+
117+
/// UID identifying a single session for the given user.
118+
var session_id: String
119+
120+
/// Used purely in language server to cancel in flight requests.
121+
/// If request_id is 0, then the request is not cancelable.
122+
/// This should be a strictly monotonically increasing number
123+
/// for the duration of a session.
124+
var request_id: UInt64
125+
}
126+
127+
enum CodeiumState: String, Codable {
128+
case unspecified = "CODEIUM_STATE_UNSPECIFIED"
129+
case inactive = "CODEIUM_STATE_INACTIVE"
130+
case processing = "CODEIUM_STATE_PROCESSING"
131+
case success = "CODEIUM_STATE_SUCCESS"
132+
case warning = "CODEIUM_STATE_WARNING"
133+
case error = "CODEIUM_STATE_ERROR"
134+
}
135+
136+
struct State: Codable {
137+
var state: CodeiumState
138+
var message: String
139+
}
140+

‎Core/Sources/CodeiumService/CodeiumRequest.swift‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
import Foundation
2+
import JSONRPC
3+
import LanguageServerProtocol
4+
import SuggestionModel
5+
6+
protocol CodeiumRequestType {
7+
associatedtype Response: Codable
8+
func makeURLRequest(server: String) -> URLRequest
9+
}
10+
11+
struct CodeiumResponseError: Codable, Error, LocalizedError {
12+
var code: String
13+
var message: String
14+
var errorDescription: String? { message }
15+
}
16+
17+
enum CodeiumRequest {
18+
struct GetCompletion: CodeiumRequestType {
19+
struct Response: Codable {
20+
var state: State
21+
var completionItems: [CodeiumCompletionItem]?
22+
}
23+
24+
struct Request: Codable {
25+
var metadata: Metadata
26+
var document: CodeiumDocument
27+
var editor_options: CodeiumEditorOptions
28+
var other_documents: [CodeiumDocument]
29+
}
30+
31+
var requestBody: Request
32+
33+
func makeURLRequest(server: String) -> URLRequest {
34+
var request = URLRequest(url: .init(string: "\(server)/exa.language_server_pb.LanguageServerService/GetCompletions")!)
35+
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
36+
let data = (try? JSONEncoder().encode(requestBody)) ?? Data() //
37+
38+
request.httpMethod = "POST"
39+
request.httpBody = data
40+
41+
return request
42+
}
43+
44+
}
45+
}
46+

0 commit comments

Comments
 (0)