Skip to content

Commit

Permalink
Add code execution support
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard committed Jul 29, 2024
1 parent f012e91 commit e96ed04
Show file tree
Hide file tree
Showing 4 changed files with 123 additions and 5 deletions.
3 changes: 2 additions & 1 deletion Sources/GoogleAI/Chat.swift
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,8 @@ public class Chat {
case let .text(str):
combinedText += str

case .data, .fileData, .functionCall, .functionResponse:
case .data, .fileData, .functionCall, .functionResponse, .executableCode,
.codeExecutionResult:
// Don't combine it, just add to the content. If there's any text pending, add that as
// a part.
if !combinedText.isEmpty {
Expand Down
85 changes: 84 additions & 1 deletion Sources/GoogleAI/FunctionCalling.swift
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,9 @@ public struct Tool {
/// A list of `FunctionDeclarations` available to the model.
let functionDeclarations: [FunctionDeclaration]?

/// Enables the model to execute code as part of generation.
let codeExecution: CodeExecution?

/// Constructs a new `Tool`.
///
/// - Parameters:
Expand All @@ -172,8 +175,11 @@ public struct Tool {
/// populating ``FunctionCall`` in the response. The next conversation turn may contain a
/// ``FunctionResponse`` in ``ModelContent/Part/functionResponse(_:)`` with the
/// ``ModelContent/role`` "function", providing generation context for the next model turn.
public init(functionDeclarations: [FunctionDeclaration]?) {
/// - codeExecution: Enables the model to execute code as part of generation, if provided.
public init(functionDeclarations: [FunctionDeclaration]? = nil,
codeExecution: CodeExecution? = nil) {
self.functionDeclarations = functionDeclarations
self.codeExecution = codeExecution
}
}

Expand Down Expand Up @@ -244,6 +250,55 @@ public struct FunctionResponse: Equatable {
}
}

/// Tool that executes code generated by the model, automatically returning the result to the model.
///
/// This type has no fields. See ``ExecutableCode`` and ``CodeExecutionResult``, which are only
/// generated when using this tool.
public struct CodeExecution {
/// Constructs a new `CodeExecution` tool.
public init() {}
}

/// Code generated by the model that is meant to be executed, and the result returned to the model.
///
/// Only generated when using the ``CodeExecution`` tool, in which case the code will automatically
/// be executed, and a corresponding ``CodeExecutionResult`` will also be generated.
public struct ExecutableCode: Equatable {
/// The programming language of the ``code``.
public let language: String

/// The code to be executed.
public let code: String
}

/// Result of executing the ``ExecutableCode``.
///
/// Only generated when using the ``CodeExecution`` tool, and always follows a part containing the
/// ``ExecutableCode``.
public struct CodeExecutionResult: Equatable {
/// Possible outcomes of the code execution.
public enum Outcome: String {
/// An unrecognized code execution outcome was provided.
case unknown = "OUTCOME_UNKNOWN"
/// Unspecified status; this value should not be used.
case unspecified = "OUTCOME_UNSPECIFIED"
/// Code execution completed successfully.
case ok = "OUTCOME_OK"
/// Code execution finished but with a failure; ``CodeExecutionResult/output`` should contain
/// the failure details from `stderr`.
case failed = "OUTCOME_FAILED"
/// Code execution ran for too long, and was cancelled. There may or may not be a partial
/// ``CodeExecutionResult/output`` present.
case deadlineExceeded = "OUTCOME_DEADLINE_EXCEEDED"
}

/// Outcome of the code execution.
public let outcome: Outcome

/// Contains `stdout` when code execution is successful, `stderr` or other description otherwise.
public let output: String
}

// MARK: - Codable Conformance

extension FunctionCall: Decodable {
Expand Down Expand Up @@ -293,3 +348,31 @@ extension FunctionCallingConfig.Mode: Encodable {}
extension ToolConfig: Encodable {}

extension FunctionResponse: Encodable {}

extension CodeExecution: Encodable {}

extension ExecutableCode: Codable {}

@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
extension CodeExecutionResult.Outcome: Codable {
public init(from decoder: any Decoder) throws {
let value = try decoder.singleValueContainer().decode(String.self)
guard let decodedOutcome = CodeExecutionResult.Outcome(rawValue: value) else {
Logging.default
.error("[GoogleGenerativeAI] Unrecognized Outcome with value \"\(value)\".")
self = .unknown
return
}

self = decodedOutcome
}
}

@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
extension CodeExecutionResult: Codable {
public init(from decoder: any Decoder) throws {
let container = try decoder.container(keyedBy: CodingKeys.self)
outcome = try container.decode(Outcome.self, forKey: .outcome)
output = try container.decodeIfPresent(String.self, forKey: .output) ?? ""
}
}
21 changes: 18 additions & 3 deletions Sources/GoogleAI/GenerateContentResponse.swift
Original file line number Diff line number Diff line change
Expand Up @@ -46,16 +46,31 @@ public struct GenerateContentResponse {
return nil
}
let textValues: [String] = candidate.content.parts.compactMap { part in
guard case let .text(text) = part else {
switch part {
case let .text(text):
return text
case let .executableCode(executableCode):
let codeBlockLanguage: String
if executableCode.language == "LANGUAGE_UNSPECIFIED" {
codeBlockLanguage = ""
} else {
codeBlockLanguage = executableCode.language.lowercased()
}
return "```\(codeBlockLanguage)\n\(executableCode.code)\n```"
case let .codeExecutionResult(codeExecutionResult):
if codeExecutionResult.output.isEmpty {
return nil
}
return "```\n\(codeExecutionResult.output)\n```"
case .data, .fileData, .functionCall, .functionResponse:
return nil
}
return text
}
guard textValues.count > 0 else {
Logging.default.error("Could not get a text part from the first candidate.")
return nil
}
return textValues.joined(separator: " ")
return textValues.joined(separator: "\n")
}

/// Returns function calls found in any `Part`s of the first candidate of the response, if any.
Expand Down
19 changes: 19 additions & 0 deletions Sources/GoogleAI/ModelContent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,12 @@ public struct ModelContent: Equatable {
/// A response to a function call.
case functionResponse(FunctionResponse)

/// Code generated by the model that is meant to be executed.
case executableCode(ExecutableCode)

/// Result of executing the ``ExecutableCode``.
case codeExecutionResult(CodeExecutionResult)

// MARK: Convenience Initializers

/// Convenience function for populating a Part with JPEG data.
Expand Down Expand Up @@ -129,6 +135,8 @@ extension ModelContent.Part: Codable {
case fileData
case functionCall
case functionResponse
case executableCode
case codeExecutionResult
}

enum InlineDataKeys: String, CodingKey {
Expand Down Expand Up @@ -164,6 +172,10 @@ extension ModelContent.Part: Codable {
try container.encode(functionCall, forKey: .functionCall)
case let .functionResponse(functionResponse):
try container.encode(functionResponse, forKey: .functionResponse)
case let .executableCode(executableCode):
try container.encode(executableCode, forKey: .executableCode)
case let .codeExecutionResult(codeExecutionResult):
try container.encode(codeExecutionResult, forKey: .codeExecutionResult)
}
}

Expand All @@ -181,6 +193,13 @@ extension ModelContent.Part: Codable {
self = .data(mimetype: mimetype, bytes)
} else if values.contains(.functionCall) {
self = try .functionCall(values.decode(FunctionCall.self, forKey: .functionCall))
} else if values.contains(.executableCode) {
self = try .executableCode(values.decode(ExecutableCode.self, forKey: .executableCode))
} else if values.contains(.codeExecutionResult) {
self = try .codeExecutionResult(values.decode(
CodeExecutionResult.self,
forKey: .codeExecutionResult
))
} else {
throw DecodingError.dataCorrupted(.init(
codingPath: [CodingKeys.text, CodingKeys.inlineData],
Expand Down

0 comments on commit e96ed04

Please sign in to comment.