introduced new auto function declaration type and auto function calling

Author: David Motsonashvili

firebase-ai/src/main/kotlin/com/google/firebase/ai/Chat.kt

@@ -18,17 +18,18 @@ package com.google.firebase.ai
 
 import android.graphics.Bitmap
 import com.google.firebase.ai.type.Content
+import com.google.firebase.ai.type.FunctionCallPart
+import com.google.firebase.ai.type.FunctionResponsePart
 import com.google.firebase.ai.type.GenerateContentResponse
-import com.google.firebase.ai.type.ImagePart
-import com.google.firebase.ai.type.InlineDataPart
 import com.google.firebase.ai.type.InvalidStateException
 import com.google.firebase.ai.type.TextPart
 import com.google.firebase.ai.type.content
 import java.util.LinkedList
 import java.util.concurrent.Semaphore
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.FlowCollector
 import kotlinx.coroutines.flow.onCompletion
-import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.transform
 
 /**
  * Representation of a multi-turn interaction with a model.
@@ -51,25 +52,36 @@ public class Chat(
   private var lock = Semaphore(1)
 
   /**
-   * Sends a message using the provided [prompt]; automatically providing the existing [history] as
-   * context.
+   * Sends a message using the provided [inputPrompt]; automatically providing the existing
+   * [history] as context.
    *
    * If successful, the message and response will be added to the [history]. If unsuccessful,
    * [history] will remain unchanged.
    *
-   * @param prompt The input that, together with the history, will be given to the model as the
+   * @param inputPrompt The input that, together with the history, will be given to the model as the
    * prompt.
-   * @throws InvalidStateException if [prompt] is not coming from the 'user' role.
+   * @throws InvalidStateException if [inputPrompt] is not coming from the 'user' role.
    * @throws InvalidStateException if the [Chat] instance has an active request.
    */
-  public suspend fun sendMessage(prompt: Content): GenerateContentResponse {
-    prompt.assertComesFromUser()
+  public suspend fun sendMessage(inputPrompt: Content): GenerateContentResponse {
+    inputPrompt.assertComesFromUser()
     attemptLock()
+    var response: GenerateContentResponse
+    var prompt = inputPrompt
     try {
-      val fullPrompt = history + prompt
-      val response = model.generateContent(fullPrompt.first(), *fullPrompt.drop(1).toTypedArray())
-      history.add(prompt)
-      history.add(response.candidates.first().content)
+      while (true) {
+        response = model.generateContent(listOf(*history.toTypedArray(), prompt))
+        val responsePart = response.candidates.first().content.parts.first()
+
+        history.add(prompt)
+        history.add(response.candidates.first().content)
+        if (responsePart is FunctionCallPart) {
+          val output = model.executeFunction(responsePart)
+          prompt = Content("function", listOf(FunctionResponsePart(responsePart.name, output)))
+        } else {
+          break
+        }
+      }
       return response
     } finally {
       lock.release()
@@ -130,43 +142,20 @@ public class Chat(
 
     val fullPrompt = history + prompt
     val flow = model.generateContentStream(fullPrompt.first(), *fullPrompt.drop(1).toTypedArray())
-    val bitmaps = LinkedList<Bitmap>()
-    val inlineDataParts = LinkedList<InlineDataPart>()
-    val text = StringBuilder()
+    val tempHistory = LinkedList<Content>()
+    tempHistory.add(prompt)
 
     /**
      * TODO: revisit when images and inline data are returned. This will cause issues with how
      * things are structured in the response. eg; a text/image/text response will be (incorrectly)
      * represented as image/text
      */
     return flow
-      .onEach {
-        for (part in it.candidates.first().content.parts) {
-          when (part) {
-            is TextPart -> text.append(part.text)
-            is ImagePart -> bitmaps.add(part.image)
-            is InlineDataPart -> inlineDataParts.add(part)
-          }
-        }
-      }
+      .transform { response -> automaticFunctionExecutingTransform(this, tempHistory, response) }
       .onCompletion {
         lock.release()
         if (it == null) {
-          val content =
-            content("model") {
-              for (bitmap in bitmaps) {
-                image(bitmap)
-              }
-              for (inlineDataPart in inlineDataParts) {
-                inlineData(inlineDataPart.inlineData, inlineDataPart.mimeType)
-              }
-              if (text.isNotBlank()) {
-                text(text.toString())
-              }
-            }
-
-          history.add(prompt)
-          history.add(content)
+          history.addAll(tempHistory)
         }
       }
   }
@@ -209,6 +198,62 @@ public class Chat(
     return sendMessageStream(content)
   }
 
+  private suspend fun automaticFunctionExecutingTransform(
+    transformer: FlowCollector<GenerateContentResponse>,
+    tempHistory: LinkedList<Content>,
+    response: GenerateContentResponse
+  ) {
+    for (part in response.candidates.first().content.parts) {
+      when (part) {
+        is TextPart -> {
+          transformer.emit(response)
+          addTextToHistory(tempHistory, part)
+        }
+        is FunctionCallPart -> {
+          val functionCall =
+            response.candidates.first().content.parts.first { it is FunctionCallPart }
+              as FunctionCallPart
+          val output = model.executeFunction(functionCall)
+          val functionResponse =
+            Content("function", listOf(FunctionResponsePart(functionCall.name, output)))
+          tempHistory.add(response.candidates.first().content)
+          tempHistory.add(functionResponse)
+          model
+            .generateContentStream(listOf(*history.toTypedArray(), *tempHistory.toTypedArray()))
+            .collect { automaticFunctionExecutingTransform(transformer, tempHistory, it) }
+        }
+        else -> {
+          transformer.emit(response)
+          tempHistory.add(Content("model", listOf(part)))
+        }
+      }
+    }
+  }
+
+  private fun addTextToHistory(tempHistory: LinkedList<Content>, textPart: TextPart) {
+    val lastContent = tempHistory.lastOrNull()
+    if (lastContent?.role == "model" && lastContent.parts.any { it is TextPart }) {
+      tempHistory.removeLast()
+      val editedContent =
+        Content(
+          "model",
+          lastContent.parts.map {
+            when (it) {
+              is TextPart -> {
+                TextPart(it.text + textPart.text)
+              }
+              else -> {
+                it
+              }
+            }
+          }
+        )
+      tempHistory.add(editedContent)
+      return
+    }
+    tempHistory.add(Content("model", listOf(textPart)))
+  }
+
   private fun Content.assertComesFromUser() {
     if (role !in listOf("user", "function")) {
       throw InvalidStateException("Chat prompts should come from the 'user' or 'function' role.")

firebase-ai/src/main/kotlin/com/google/firebase/ai/GenerativeModel.kt

@@ -22,10 +22,12 @@ import com.google.firebase.ai.common.APIController
 import com.google.firebase.ai.common.AppCheckHeaderProvider
 import com.google.firebase.ai.common.CountTokensRequest
 import com.google.firebase.ai.common.GenerateContentRequest
+import com.google.firebase.ai.type.AutoFunctionDeclaration
 import com.google.firebase.ai.type.Content
 import com.google.firebase.ai.type.CountTokensResponse
 import com.google.firebase.ai.type.FinishReason
 import com.google.firebase.ai.type.FirebaseAIException
+import com.google.firebase.ai.type.FunctionCallPart
 import com.google.firebase.ai.type.GenerateContentResponse
 import com.google.firebase.ai.type.GenerationConfig
 import com.google.firebase.ai.type.GenerativeBackend
@@ -45,6 +47,11 @@ import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.map
 import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.InternalSerializationApi
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.serializerOrNull
 
 /**
  * Represents a multimodal model (like Gemini), capable of generating content based on various input
@@ -266,6 +273,43 @@ internal constructor(
     return countTokens(content { image(prompt) })
   }
 
+  @OptIn(InternalSerializationApi::class)
+  internal suspend fun executeFunction(call: FunctionCallPart): JsonObject {
+    if (tools == null) {
+      throw RuntimeException("No registered tools")
+    }
+    val tool = tools.flatMap { it.autoFunctionDeclarations?.filterNotNull() ?: emptyList() }
+    val declaration =
+      tool.firstOrNull() { it.name == call.name }
+        ?: throw RuntimeException("No registered function named ${call.name}")
+    return executeFunction<Any, Any>(
+      declaration as AutoFunctionDeclaration<Any, Any>,
+      call.args["param"].toString()
+    )
+  }
+
+  @OptIn(InternalSerializationApi::class)
+  internal suspend fun <I : Any, O : Any> executeFunction(
+    functionDeclaration: AutoFunctionDeclaration<I, O>,
+    parameter: String
+  ): JsonObject {
+    val inputDeserializer =
+      functionDeclaration.inputSchema.clazz.serializerOrNull()
+        ?: throw RuntimeException(
+          "Function input type ${functionDeclaration.inputSchema.clazz.qualifiedName} is not @Serializable"
+        )
+    val input = Json.decodeFromString(inputDeserializer, parameter)
+    val functionReference =
+      functionDeclaration.functionReference
+        ?: throw RuntimeException("Function reference for ${functionDeclaration.name} is missing")
+    val output = functionReference.invoke(input)
+    val outputSerializer = functionDeclaration.outputSchema?.clazz?.serializerOrNull()
+    if (outputSerializer != null) {
+      return Json.encodeToJsonElement(outputSerializer, output).jsonObject
+    }
+    return output as JsonObject
+  }
+
   @OptIn(ExperimentalSerializationApi::class)
   private fun constructRequest(vararg prompt: Content) =
     GenerateContentRequest(

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/AutoFunctionDeclaration.kt

@@ -0,0 +1,55 @@
+package com.google.firebase.ai.type
+
+import kotlinx.serialization.json.JsonObject
+
+public class AutoFunctionDeclaration<I : Any, O : Any>
+internal constructor(
+  public val name: String,
+  public val description: String,
+  public val inputSchema: JsonSchema<I>,
+  public val outputSchema: JsonSchema<O>? = null,
+  public val functionReference: (suspend (I) -> O)? = null
+) {
+  public companion object {
+    public fun <I : Any, O : Any> create(
+      functionName: String,
+      description: String,
+      inputSchema: JsonSchema<I>,
+      outputSchema: JsonSchema<O>,
+      functionReference: ((I) -> O)? = null
+    ): AutoFunctionDeclaration<I, O> {
+      return AutoFunctionDeclaration<I, O>(
+        functionName,
+        description,
+        inputSchema,
+        outputSchema,
+        functionReference
+      )
+    }
+
+    public fun <I : Any> create(
+      functionName: String,
+      inputSchema: JsonSchema<I>,
+      description: String,
+      functionReference: ((I) -> JsonObject)? = null
+    ): AutoFunctionDeclaration<I, JsonObject> {
+      return AutoFunctionDeclaration<I, JsonObject>(
+        functionName,
+        description,
+        inputSchema,
+        null,
+        functionReference
+      )
+    }
+  }
+
+  internal fun toInternal(): FunctionDeclaration.Internal {
+    return FunctionDeclaration.Internal(
+      name,
+      description,
+      null,
+      JsonSchema.obj(mapOf("param" to inputSchema)).toInternalJson(),
+      outputSchema?.toInternalJson()
+    )
+  }
+}

firebase-ai/src/main/kotlin/com/google/firebase/ai/type/FunctionDeclaration.kt

@@ -61,12 +61,14 @@ public class FunctionDeclaration(
   internal val schema: Schema =
     Schema.obj(properties = parameters, optionalProperties = optionalParameters, nullable = false)
 
-  internal fun toInternal() = Internal(name, description, schema.toInternalOpenApi())
+  internal fun toInternal() = Internal(name, description, schema.toInternalOpenApi(), null, null)
 
   @Serializable
   internal data class Internal(
     val name: String,
     val description: String,
-    val parameters: Schema.InternalOpenAPI
+    val parameters: Schema.InternalOpenAPI?,
+    val parametersJsonSchema: Schema.InternalJson?,
+    val responseJsonSchema: Schema.InternalJson?,
   )
 }