Skip to content

Commit

Permalink
Upload file testing
Browse files Browse the repository at this point in the history
  • Loading branch information
@Archetapp
Archetapp committed Dec 12, 2024
1 parent 99cdef6 commit e26b80a
Show file tree
Hide file tree
Showing 7 changed files with 255 additions and 155 deletions.
151 changes: 75 additions & 76 deletions Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.RequestBodies.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,6 @@ import Foundation

extension _Gemini.APISpecification {
public enum RequestBodies {

public struct InitiateUploadInput: Codable {
public let file: FileMetadata
public let contentLength: Int
public let mimeType: String

public struct FileMetadata: Codable {
public let displayName: String

private enum CodingKeys: String, CodingKey {
case displayName = "display_name"
}
}
}

public struct CompleteUploadInput: Codable {
public let uploadURL: URL
public let fileData: Data
Expand All @@ -34,6 +19,7 @@ extension _Gemini.APISpecification {
self.offset = offset
}
}

public struct GenerateContentInput: Codable {
public let model: String
public let requestBody: SpeechRequest
Expand Down Expand Up @@ -102,6 +88,28 @@ extension _Gemini.APISpecification {
case mimeType
}

public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
switch self {
case .text(let txt):
try container.encode(txt, forKey: .text)
case .inline(data: let data, mimeType: let mimeType):
var nestedContainer = container.nestedContainer(
keyedBy: InlineDataNestedKeys.self,
forKey: .inlineData
)
try nestedContainer.encode(data.base64EncodedString(), forKey: .data)
try nestedContainer.encode(mimeType, forKey: .mimeType)
case .file(url: let url, mimeType: let mimeType):
var nestedContainer = container.nestedContainer(
keyedBy: FileDataNestedKeys.self,
forKey: .fileData
)
try nestedContainer.encode(url.absoluteString, forKey: .fileUri)
try nestedContainer.encode(mimeType, forKey: .mimeType)
}
}

public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)

Expand All @@ -120,10 +128,12 @@ extension _Gemini.APISpecification {
}

if let fileContainer = try? container.nestedContainer(keyedBy: FileDataNestedKeys.self, forKey: .fileData) {
let url = try fileContainer.decode(URL.self, forKey: .fileUri)
let urlString = try fileContainer.decode(String.self, forKey: .fileUri)
let mimeType = try fileContainer.decode(String.self, forKey: .mimeType)
self = .file(url: url, mimeType: mimeType)
return
if let url = URL(string: urlString) {
self = .file(url: url, mimeType: mimeType)
return
}
}

throw DecodingError.dataCorrupted(
Expand All @@ -133,95 +143,80 @@ extension _Gemini.APISpecification {
)
)
}

public func encode(to encoder: Encoder) throws {
var container = encoder.container(keyedBy: CodingKeys.self)
switch self {
case .text(let txt):
try container.encode(txt, forKey: .text)
case .inline(data: let data, mimeType: let mimeType):
var nestedContainer = container.nestedContainer(
keyedBy: InlineDataNestedKeys.self,
forKey: .inlineData
)
try nestedContainer.encode(data.base64EncodedString(), forKey: .data)
try nestedContainer.encode(mimeType, forKey: .mimeType)
case .file(url: let url, mimeType: let mimeType):
var nestedContainer = container.nestedContainer(
keyedBy: FileDataNestedKeys.self,
forKey: .fileData
)
try nestedContainer.encode(url, forKey: .fileUri)
try nestedContainer.encode(mimeType, forKey: .mimeType)
}
}
}
}

public struct GenerationConfig: Codable {
public let maxTokens: Int?
public let maxOutputTokens: Int?
public let temperature: Double?
public let topP: Double?
public let topK: Int?
public let presencePenalty: Double?
public let frequencyPenalty: Double?
public let responseMimeType: String?

public init(
maxTokens: Int? = nil,
maxOutputTokens: Int? = nil,
temperature: Double? = nil,
topP: Double? = nil,
topK: Int? = nil,
presencePenalty: Double? = nil,
frequencyPenalty: Double? = nil
frequencyPenalty: Double? = nil,
responseMimeType: String? = nil
) {
self.maxTokens = maxTokens
self.maxOutputTokens = maxOutputTokens
self.temperature = temperature
self.topP = topP
self.topK = topK
self.presencePenalty = presencePenalty
self.frequencyPenalty = frequencyPenalty
self.responseMimeType = responseMimeType
}
}

public struct FileUploadInput: Codable, HTTPRequest.Multipart.ContentConvertible {
public struct FileUploadInput: Codable {
public let fileData: Data
public let mimeType: String
public let model: String
public let displayName: String

public init(
fileData: Data,
mimeType: String,
model: _Gemini.Model
) {
public struct Metadata: Codable {
public let file: File

public struct File: Codable {
let displayName: String

private enum CodingKeys: String, CodingKey {
case displayName = "display_name"
}
}
}

public init(fileData: Data, mimeType: String, displayName: String) {
self.fileData = fileData
self.mimeType = mimeType
self.model = model.rawValue
self.displayName = displayName
}

public func __conversion() throws -> HTTPRequest.Multipart.Content {
var result = HTTPRequest.Multipart.Content()

// Add the JSON part
let jsonBody = ["file": ["mimeType": self.mimeType]]
let jsonData = try JSONSerialization.data(withJSONObject: jsonBody)
result.append(
.text(
named: "metadata",
value: String(data: jsonData, encoding: .utf8) ?? ""
)
)

// Add the file part
result.append(
.file(
named: "file",
data: fileData,
filename: "file",
contentType: .init(rawValue: mimeType)
)
)

return result
private enum CodingKeys: String, CodingKey {
case file
case fileData
case mimeType
case displayName
}

public func encode(to encoder: Encoder) throws {
// Encode only the metadata part as JSON
var container = encoder.container(keyedBy: CodingKeys.self)
let metadata = Metadata(file: .init(displayName: displayName))
try container.encode(metadata.file, forKey: .file)
}

public init(from decoder: Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
let metadata = try container.decode(Metadata.self, forKey: .file)
self.displayName = metadata.file.displayName
self.fileData = try container.decode(Data.self, forKey: .fileData)
self.mimeType = try container.decode(String.self, forKey: .mimeType)
}
}

Expand All @@ -232,5 +227,9 @@ extension _Gemini.APISpecification {
self.fileURL = fileURL
}
}

public struct FileStatusInput: Codable {
public let name: String
}
}
}
46 changes: 17 additions & 29 deletions Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.swift
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,22 +76,17 @@ extension _Gemini {
@POST
@Path("/upload/v1beta/files")
@Header([
"X-Goog-Upload-Protocol": "resumable",
"X-Goog-Upload-Command": "start",
"Content-Type": "application/json"
"X-Goog-Upload-Command": "start, upload, finalize"
])
var initiateUpload = Endpoint<RequestBodies.InitiateUploadInput, ResponseBodies.UploadInitiation, Void>()
@Body(json: .input)
var uploadFile = Endpoint<RequestBodies.FileUploadInput, ResponseBodies.FileUpload, Void>()

// Complete Upload endpoint
@POST
@Header([
"X-Goog-Upload-Command": "upload, finalize"
])
// File Status endpoint
@GET
@Path({ context -> String in
context.input.uploadURL.absoluteString
"/v1beta/\(context.input.name)"
})
@Body(json: \.input)
var completeUpload = Endpoint<RequestBodies.CompleteUploadInput, ResponseBodies.FileUpload, Void>()
var getFileStatus = Endpoint<RequestBodies.FileStatusInput, _Gemini.File, Void>()

// Delete File endpoint
@DELETE
Expand All @@ -108,42 +103,35 @@ extension _Gemini.APISpecification {
from input: Input,
context: BuildRequestContext
) throws -> Request {

var request = try super.buildRequestBase(
from: input,
context: context
)

if let apiKey = context.root.configuration.apiKey {
request = request.header("Authorization", "Bearer \(apiKey)")
request = request.query([.init(name: "key", value: apiKey)])
}

if let uploadInput = input as? RequestBodies.InitiateUploadInput {
request = request.header("X-Goog-Upload-Header-Content-Length", "\(uploadInput.contentLength)")
request = request.header("X-Goog-Upload-Header-Content-Type", uploadInput.mimeType)
if let uploadInput = input as? RequestBodies.FileUploadInput {
request = request.header("Content-Type", uploadInput.mimeType)
let metadata = ["file": ["display_name": uploadInput.displayName]]
let jsonData = try JSONSerialization.data(withJSONObject: metadata)

var combinedData = jsonData
combinedData.append(uploadInput.fileData)

request = request.body(.data(combinedData))
}

if let completeInput = input as? RequestBodies.CompleteUploadInput {
request = request.header("Content-Length", "\(completeInput.fileData.count)")
request = request.header("X-Goog-Upload-Offset", "\(completeInput.offset)")
}
print(request)

return request
}

override public func decodeOutputBase(
from response: Request.Response,
context: DecodeOutputContext
) throws -> Output {

print(response)
try response.validate()

if Output.self == Data.self {
return response.data as! Output
}

return try response.decode(
Output.self,
keyDecodingStrategy: .convertFromSnakeCase
Expand Down
7 changes: 7 additions & 0 deletions Sources/_Gemini/Intramodular/Models/_Gemini.MediaType.swift
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
//
// _Gemini.MediaType.swift
// AI
//
// Created by Jared Davidson on 12/11/24.
//

Loading

0 comments on commit e26b80a

Please sign in to comment.