Add code execution support

Author: andrewheard

Sources/GoogleAI/Chat.swift

@@ -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 {

Sources/GoogleAI/FunctionCalling.swift

@@ -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:
@@ -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
   }
 }
 
@@ -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 {
@@ -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) ?? ""
+  }
+}

Sources/GoogleAI/GenerateContentResponse.swift

@@ -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.

Sources/GoogleAI/ModelContent.swift

@@ -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.
@@ -129,6 +135,8 @@ extension ModelContent.Part: Codable {
     case fileData
     case functionCall
     case functionResponse
+    case executableCode
+    case codeExecutionResult
   }
 
   enum InlineDataKeys: String, CodingKey {
@@ -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)
     }
   }
 
@@ -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],