Merge pull request #2988 from wordpress-mobile/ai-rest-client-add-params

New REST AI client parameters

Author: JorgeMucientes

fluxc/src/main/java/org/wordpress/android/fluxc/network/rest/wpcom/jetpackai/JetpackAIRestClient.kt

@@ -14,6 +14,7 @@ import org.wordpress.android.fluxc.network.rest.wpcom.WPComGsonRequest.WPComGson
 import org.wordpress.android.fluxc.network.rest.wpcom.WPComGsonRequestBuilder
 import org.wordpress.android.fluxc.network.rest.wpcom.WPComGsonRequestBuilder.Response
 import org.wordpress.android.fluxc.network.rest.wpcom.auth.AccessToken
+import org.wordpress.android.fluxc.utils.extensions.putIfNotNull
 import javax.inject.Inject
 import javax.inject.Named
 import javax.inject.Singleton
@@ -55,7 +56,9 @@ class JetpackAIRestClient @Inject constructor(
     suspend fun fetchJetpackAITextCompletion(
         token: JWTToken,
         prompt: String,
-        feature: String
+        feature: String,
+        format: ResponseFormat? = null,
+        model: String? = null
     ): JetpackAICompletionsResponse {
         val url = WPCOMV2.text_completion.url
         val body = mutableMapOf<String, String>()
@@ -64,6 +67,7 @@ class JetpackAIRestClient @Inject constructor(
             put("prompt", prompt)
             put("feature", feature)
             put("_fields", FIELDS_TO_REQUEST)
+            putIfNotNull("response_format" to format?.value, "model" to model)
         }
 
         val response = wpComGsonRequestBuilder.syncPostRequest(
@@ -161,6 +165,11 @@ class JetpackAIRestClient @Inject constructor(
         NETWORK_ERROR
     }
 
+    enum class ResponseFormat(val value: String) {
+        JSON("json_object"),
+        TEXT("text")
+    }
+
     private fun WPComGsonNetworkError.toJetpackAICompletionsError() =
         when (type) {
             GenericErrorType.TIMEOUT -> JetpackAICompletionsErrorType.TIMEOUT

fluxc/src/main/java/org/wordpress/android/fluxc/store/jetpackai/JetpackAIStore.kt

@@ -8,6 +8,7 @@ import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestCli
 import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient.JetpackAIJWTTokenResponse
 import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient.JetpackAIJWTTokenResponse.Error
 import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient.JetpackAIJWTTokenResponse.Success
+import org.wordpress.android.fluxc.network.rest.wpcom.jetpackai.JetpackAIRestClient.ResponseFormat
 import org.wordpress.android.fluxc.tools.CoroutineEngine
 import org.wordpress.android.util.AppLog
 import javax.inject.Inject
@@ -70,7 +71,9 @@ class JetpackAIStore @Inject constructor(
     suspend fun fetchJetpackAICompletions(
         site: SiteModel,
         prompt: String,
-        feature: String
+        feature: String,
+        responseFormat: ResponseFormat? = null,
+        model: String? = null
     ): JetpackAICompletionsResponse = coroutineEngine.withDefaultContext(
         tag = AppLog.T.API,
         caller = this,
@@ -93,14 +96,20 @@ class JetpackAIStore @Inject constructor(
                 }
             }
 
-        val result = jetpackAIRestClient.fetchJetpackAITextCompletion(token, prompt, feature)
+        val result = jetpackAIRestClient.fetchJetpackAITextCompletion(
+            token,
+            prompt,
+            feature,
+            responseFormat,
+            model
+        )
 
         return@withDefaultContext when {
             // Fetch token anew if using existing token returns AUTH_ERROR
             result is JetpackAICompletionsResponse.Error && result.type == AUTH_ERROR -> {
                 // Remove cached token
                 this@JetpackAIStore.token = null
-                fetchJetpackAICompletions(site, prompt, feature)
+                fetchJetpackAICompletions(site, prompt, feature, responseFormat, model)
             }
 
             else -> result