Google AI: symmetric decorate() decorator

README.md

@@ -178,6 +178,27 @@ interpreter.run(inputBuffer, outputBuffer)
 interpreter.close()
 ```
 
+### Google AI (Gemini)
+
+```kotlin
+val gemini = wildEdge.decorate(
+    GenerativeModel(modelName = "gemini-2.0-flash", apiKey = "<key>"),
+    modelId = "gemini-2.0-flash",
+    modelFamily = "gemini",
+)
+
+// Streaming — tracking fires when the flow completes
+gemini.generateContentStream(prompt, inputMeta = WildEdge.analyzeText(prompt))
+    .collect { response -> append(response.text.orEmpty()) }
+
+// Unary
+val response = gemini.generateContent(prompt, inputMeta = WildEdge.analyzeText(prompt))
+
+// Untracked operations go through .model directly
+val chat = gemini.model.startChat(history)
+val tokens = gemini.model.countTokens(prompt)
+```
+
 ### Remote models
 
 ```kotlin
@@ -361,6 +382,10 @@ Integrate the WildEdge Android SDK (dev.wildedge:wildedge-android) into this pro
    - ONNX: val session = wildEdge.decorate(env.createSession(...), modelFile, modelVersion = "...")
    - LiteRT: val engine = wildEdge.decorate(Engine(config), config, modelVersion = "...")
    - MLKit: wildEdge.registerMlKitModel(...) and Task.trackWith(handle)
+   - Google AI (Gemini): val gemini = wildEdge.decorate(GenerativeModel(modelName = "...", apiKey = "..."), modelId = "gemini-2.0-flash")
+     Use gemini.generateContentStream(prompt, inputMeta = WildEdge.analyzeText(prompt)) for streaming,
+     gemini.generateContent(prompt, inputMeta = WildEdge.analyzeText(prompt)) for unary.
+     Untracked operations (startChat, countTokens) go through gemini.model directly.
    - Remote LLM: wildEdge.registerModel("id", ModelInfo(
          inputModality = InputModality.Text, outputModality = OutputModality.Generation, ...))
      Use handle.trackSuspendInference { } for suspend calls. Pass outputMetaExtractor to

examples/googleai/GoogleAiExample.kt

@@ -1,37 +1,33 @@
 package examples.googleai
 
 import com.google.ai.client.generativeai.GenerativeModel
+import com.google.ai.client.generativeai.type.Content
 import dev.wildedge.sdk.WildEdge
 import dev.wildedge.sdk.analysis.analyzeText
-import dev.wildedge.sdk.integrations.generateContentTracked
-import dev.wildedge.sdk.integrations.registerGoogleAiModel
-import dev.wildedge.sdk.integrations.trackWith
+import dev.wildedge.sdk.integrations.GoogleAiDecorator
+import dev.wildedge.sdk.integrations.decorate
 
 // Assumes WildEdge.init() has already run (manifest meta-data or Application.onCreate()).
 class GoogleAiExample {
 
     private val wildEdge = WildEdge.getInstance()
 
-    private val model = GenerativeModel(
-        modelName = "gemini-2.0-flash",
-        apiKey = "<YOUR_GOOGLE_AI_API_KEY>",
-    )
-
-    // One handle per model name. Register once, reuse for every call.
-    private val handle = wildEdge.registerGoogleAiModel(
+    private val gemini: GoogleAiDecorator = wildEdge.decorate(
+        GenerativeModel(
+            modelName = "gemini-2.0-flash",
+            apiKey = "<YOUR_GOOGLE_AI_API_KEY>",
+        ),
         modelId = "gemini-2.0-flash",
         modelFamily = "gemini",
     )
 
     // Streaming: tracking event fires when the flow completes.
-    fun stream(prompt: String) = model.generateContentStream(prompt)
-        .trackWith(handle, inputMeta = WildEdge.analyzeText(prompt))
+    fun stream(prompt: String) = gemini.generateContentStream(prompt, inputMeta = WildEdge.analyzeText(prompt))
 
     // Unary: tracking event fires when the suspend call returns.
-    suspend fun generate(prompt: String) = model.generateContentTracked(
-        handle = handle,
-        prompt = prompt,
-        inputMeta = WildEdge.analyzeText(prompt),
-    )
+    suspend fun generate(prompt: String) = gemini.generateContent(prompt, inputMeta = WildEdge.analyzeText(prompt))
+
+    // Untracked operations go through .model directly.
+    fun chat(history: List<Content>) = gemini.model.startChat(history)
 
 }

samples/cloud-llm/src/main/kotlin/dev/wildedge/sample/cloudllm/MainActivity.kt

@@ -9,8 +9,8 @@ import com.google.ai.client.generativeai.GenerativeModel
 import dev.wildedge.sample.cloudllm.databinding.ActivityMainBinding
 import dev.wildedge.sdk.WildEdge
 import dev.wildedge.sdk.analysis.analyzeText
-import dev.wildedge.sdk.integrations.registerGoogleAiModel
-import dev.wildedge.sdk.integrations.trackWith
+import dev.wildedge.sdk.integrations.GoogleAiDecorator
+import dev.wildedge.sdk.integrations.decorate
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
 
@@ -24,20 +24,16 @@ class MainActivity : AppCompatActivity() {
     private lateinit var binding: ActivityMainBinding
     private val wildEdge = WildEdge.getInstance()
 
-    private lateinit var model: GenerativeModel
+    private lateinit var model: GoogleAiDecorator
     private var generateJob: Job? = null
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         binding = ActivityMainBinding.inflate(layoutInflater)
         setContentView(binding.root)
 
-        model = GenerativeModel(
-            modelName = MODEL_NAME,
-            apiKey = BuildConfig.GOOGLE_AI_API_KEY,
-        )
-
-        val handle = wildEdge.registerGoogleAiModel(
+        model = wildEdge.decorate(
+            GenerativeModel(modelName = MODEL_NAME, apiKey = BuildConfig.GOOGLE_AI_API_KEY),
             modelId = MODEL_NAME,
             modelFamily = "gemini",
         )
@@ -73,8 +69,7 @@ class MainActivity : AppCompatActivity() {
 
             generateJob = lifecycleScope.launch {
                 try {
-                    model.generateContentStream(prompt)
-                        .trackWith(handle, inputMeta = inputMeta)
+                    model.generateContentStream(prompt, inputMeta = inputMeta)
                         .collect { response ->
                             val chunk = response.text.orEmpty()
                             if (chunk.isNotEmpty()) {

wildedge/build.gradle.kts

@@ -46,6 +46,7 @@ dependencies {
     compileOnly(libs.onnxruntime)
     compileOnly(libs.litertlm)
     compileOnly(libs.googleai)
+    testImplementation(libs.googleai)
 
     testImplementation(libs.junit)
     testImplementation(libs.coroutines.android)

wildedge/src/main/kotlin/dev/wildedge/sdk/integrations/GoogleAiDecorator.kt

@@ -0,0 +1,76 @@
+package dev.wildedge.sdk.integrations
+
+import android.graphics.Bitmap
+import com.google.ai.client.generativeai.GenerativeModel
+import com.google.ai.client.generativeai.type.Content
+import com.google.ai.client.generativeai.type.GenerateContentResponse
+import dev.wildedge.sdk.InputModality
+import dev.wildedge.sdk.ModelHandle
+import dev.wildedge.sdk.ModelInfo
+import dev.wildedge.sdk.OutputModality
+import dev.wildedge.sdk.WildEdgeClient
+import dev.wildedge.sdk.events.TextInputMeta
+import kotlinx.coroutines.flow.Flow
+
+/**
+ * Wraps a [GenerativeModel] to automatically record inference metrics via [ModelHandle].
+ *
+ * Create via [WildEdgeClient.decorate] rather than constructing directly.
+ * Untracked operations (e.g. [model].startChat, [model].countTokens) are available through [model].
+ */
+class GoogleAiDecorator(
+    val model: GenerativeModel,
+    val handle: ModelHandle,
+) {
+    /** Calls [GenerativeModel.generateContent] and records an inference event. */
+    suspend fun generateContent(prompt: String, inputMeta: TextInputMeta? = null): GenerateContentResponse =
+        model.generateContentTracked(handle, prompt, inputMeta)
+
+    /** Calls [GenerativeModel.generateContent] and records an inference event. */
+    suspend fun generateContent(vararg prompt: Content, inputMeta: TextInputMeta? = null): GenerateContentResponse =
+        model.generateContentTracked(handle, *prompt, inputMeta = inputMeta)
+
+    /** Calls [GenerativeModel.generateContent] and records an inference event. */
+    suspend fun generateContent(prompt: Bitmap): GenerateContentResponse =
+        model.generateContentTracked(handle, prompt)
+
+    /** Calls [GenerativeModel.generateContentStream]; records an inference event when the flow completes. */
+    fun generateContentStream(prompt: String, inputMeta: TextInputMeta? = null): Flow<GenerateContentResponse> =
+        model.generateContentStream(prompt).trackWith(handle, inputMeta)
+
+    /** Calls [GenerativeModel.generateContentStream]; records an inference event when the flow completes. */
+    fun generateContentStream(vararg prompt: Content, inputMeta: TextInputMeta? = null): Flow<GenerateContentResponse> =
+        model.generateContentStream(*prompt).trackWith(handle, inputMeta)
+
+    /** Calls [GenerativeModel.generateContentStream]; records an inference event when the flow completes. */
+    fun generateContentStream(prompt: Bitmap): Flow<GenerateContentResponse> =
+        model.generateContentStream(prompt).trackWith(handle, inputModality = InputModality.Image)
+}
+
+/**
+ * Wraps [model] in a [GoogleAiDecorator] and registers it for tracking.
+ *
+ * @param modelId Identifier for the model, e.g. `"gemini-2.0-flash"`.
+ * @param modelVersion Optional version string.
+ * @param modelFamily Model family label, defaults to `"gemini"`.
+ */
+fun WildEdgeClient.decorate(
+    model: GenerativeModel,
+    modelId: String,
+    modelVersion: String? = null,
+    modelFamily: String? = "gemini",
+): GoogleAiDecorator {
+    val handle = registerModel(
+        modelId,
+        ModelInfo(
+            modelName = modelId,
+            modelVersion = modelVersion,
+            modelSource = "api",
+            modelFormat = "api",
+            modelFamily = modelFamily,
+            inputModality = InputModality.Text,
+            outputModality = OutputModality.Generation,
+        ),
+    )
+    return GoogleAiDecorator(model, handle)
+}

wildedge/src/main/kotlin/dev/wildedge/sdk/integrations/GoogleAiTracking.kt

@@ -1,56 +1,21 @@
 package dev.wildedge.sdk.integrations
 
+import android.graphics.Bitmap
 import com.google.ai.client.generativeai.GenerativeModel
+import com.google.ai.client.generativeai.type.Content
 import com.google.ai.client.generativeai.type.GenerateContentResponse
 import dev.wildedge.sdk.InputModality
 import dev.wildedge.sdk.ModelHandle
-import dev.wildedge.sdk.ModelInfo
 import dev.wildedge.sdk.OutputModality
-import dev.wildedge.sdk.WildEdgeClient
 import dev.wildedge.sdk.analysis.approximateBpeTokenCount
 import dev.wildedge.sdk.events.GenerationOutputMeta
 import dev.wildedge.sdk.events.TextInputMeta
+import dev.wildedge.sdk.trackSuspendInference
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.flow
 
-/**
- * Registers a Google AI (Gemini) model and returns a [ModelHandle] for tracking.
- *
- * Pass a [modelId] matching the model name passed to [GenerativeModel] (e.g. `"gemini-2.0-flash"`).
- */
-fun WildEdgeClient.registerGoogleAiModel(
-    modelId: String,
-    modelName: String = modelId,
-    modelVersion: String? = null,
-    modelFamily: String? = "gemini",
-): ModelHandle = registerModel(
-    modelId,
-    ModelInfo(
-        modelName = modelName,
-        modelVersion = modelVersion,
-        modelSource = "api",
-        modelFormat = "api",
-        modelFamily = modelFamily,
-        inputModality = InputModality.Text,
-        outputModality = OutputModality.Generation,
-    ),
-)
-
-/**
- * Wraps a [Flow]<[GenerateContentResponse]> to track generation metrics.
- *
- * Token counts are read from [GenerateContentResponse.usageMetadata] on the final chunk,
- * with a char-based estimate as fallback. A tracking event is emitted on completion or error.
- *
- * Usage:
- * ```
- * model.generateContentStream(prompt)
- *     .trackWith(handle, inputMeta = WildEdge.analyzeText(prompt))
- *     .collect { response -> append(response.text.orEmpty()) }
- * ```
- */
 @Suppress("TooGenericExceptionCaught")
-fun Flow<GenerateContentResponse>.trackWith(
+internal fun Flow<GenerateContentResponse>.trackWith(
     handle: ModelHandle,
     inputMeta: TextInputMeta? = null,
     inputModality: InputModality = InputModality.Text,
@@ -66,7 +31,6 @@ fun Flow<GenerateContentResponse>.trackWith(
             val chunk = response.text.orEmpty()
             if (chunk.isNotEmpty() && firstTokenAt == null) firstTokenAt = System.currentTimeMillis()
             charCount += chunk.length
-            // usageMetadata is typically populated only on the final chunk
             response.usageMetadata?.let { meta ->
                 if (meta.promptTokenCount > 0) tokensIn = meta.promptTokenCount
                 if (meta.candidatesTokenCount > 0) tokensOut = meta.candidatesTokenCount
@@ -99,50 +63,41 @@ fun Flow<GenerateContentResponse>.trackWith(
     }
 }
 
-/**
- * Calls [GenerativeModel.generateContent] and tracks the inference event.
- *
- * Token counts are read from [GenerateContentResponse.usageMetadata].
- *
- * Usage:
- * ```
- * val response = model.generateContentTracked(
- *     handle = handle,
- *     prompt = "Summarize this article: ...",
- *     inputMeta = WildEdge.analyzeText(prompt),
- * )
- * ```
- */
-@Suppress("TooGenericExceptionCaught")
-suspend fun GenerativeModel.generateContentTracked(
+internal suspend fun GenerativeModel.generateContentTracked(
     handle: ModelHandle,
     prompt: String,
     inputMeta: TextInputMeta? = null,
-): GenerateContentResponse {
-    val start = System.currentTimeMillis()
-    return try {
-        val response = generateContent(prompt)
-        val meta = response.usageMetadata
-        handle.trackInference(
-            durationMs = (System.currentTimeMillis() - start).toInt(),
-            inputModality = InputModality.Text,
-            outputModality = OutputModality.Generation,
-            inputMeta = inputMeta?.toMap(),
-            outputMeta = GenerationOutputMeta(
-                tokensIn = meta?.promptTokenCount ?: inputMeta?.tokenCount,
-                tokensOut = meta?.candidatesTokenCount
-                    ?: approximateBpeTokenCount(response.text?.length ?: 0),
-            ).toMap(),
-        )
-        response
-    } catch (e: Exception) {
-        handle.trackInference(
-            durationMs = (System.currentTimeMillis() - start).toInt(),
-            inputModality = InputModality.Text,
-            outputModality = OutputModality.Generation,
-            success = false,
-            errorCode = e.javaClass.simpleName,
-        )
-        throw e
-    }
+): GenerateContentResponse = handle.trackSuspendInference(
+    inputModality = InputModality.Text,
+    outputModality = OutputModality.Generation,
+    inputMeta = inputMeta?.toMap(),
+    outputMetaExtractor = { r -> r.generationOutputMeta(inputMeta?.tokenCount) },
+) { generateContent(prompt) }
+
+internal suspend fun GenerativeModel.generateContentTracked(
+    handle: ModelHandle,
+    vararg prompt: Content,
+    inputMeta: TextInputMeta? = null,
+): GenerateContentResponse = handle.trackSuspendInference(
+    inputModality = InputModality.Text,
+    outputModality = OutputModality.Generation,
+    inputMeta = inputMeta?.toMap(),
+    outputMetaExtractor = { r -> r.generationOutputMeta(inputMeta?.tokenCount) },
+) { generateContent(*prompt) }
+
+internal suspend fun GenerativeModel.generateContentTracked(
+    handle: ModelHandle,
+    prompt: Bitmap,
+): GenerateContentResponse = handle.trackSuspendInference(
+    inputModality = InputModality.Image,
+    outputModality = OutputModality.Generation,
+    outputMetaExtractor = { r -> r.generationOutputMeta(tokensInFallback = null) },
+) { generateContent(prompt) }
+
+private fun GenerateContentResponse.generationOutputMeta(tokensInFallback: Int?): Map<String, Any?> {
+    val meta = usageMetadata
+    return GenerationOutputMeta(
+        tokensIn = meta?.promptTokenCount ?: tokensInFallback,
+        tokensOut = meta?.candidatesTokenCount ?: approximateBpeTokenCount(text?.length ?: 0),
+    ).toMap()
 }