Copy tests to Firebase AI

Author: emilypgoogle

firebase-ai/src/testUtil/java/com/google/firebase/ai/JavaCompileTests.java

@@ -14,14 +14,18 @@
  * limitations under the License.
  */
 
-package java.com.google.firebase.ai;
+package java.com.google.firebase.vertexai;
 
 import android.graphics.Bitmap;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.firebase.concurrent.FirebaseExecutors;
 import com.google.firebase.ai.FirebaseAI;
 import com.google.firebase.ai.GenerativeModel;
+import com.google.firebase.ai.LiveGenerativeModel;
 import com.google.firebase.ai.java.ChatFutures;
 import com.google.firebase.ai.java.GenerativeModelFutures;
+import com.google.firebase.ai.java.LiveModelFutures;
+import com.google.firebase.ai.java.LiveSessionFutures;
 import com.google.firebase.ai.type.BlockReason;
 import com.google.firebase.ai.type.Candidate;
 import com.google.firebase.ai.type.Citation;
@@ -32,25 +36,33 @@
 import com.google.firebase.ai.type.FileDataPart;
 import com.google.firebase.ai.type.FinishReason;
 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.GenerationConfig;
 import com.google.firebase.ai.type.HarmCategory;
 import com.google.firebase.ai.type.HarmProbability;
 import com.google.firebase.ai.type.HarmSeverity;
 import com.google.firebase.ai.type.ImagePart;
 import com.google.firebase.ai.type.InlineDataPart;
+import com.google.firebase.ai.type.LiveContentResponse;
+import com.google.firebase.ai.type.LiveGenerationConfig;
+import com.google.firebase.ai.type.MediaData;
 import com.google.firebase.ai.type.ModalityTokenCount;
 import com.google.firebase.ai.type.Part;
 import com.google.firebase.ai.type.PromptFeedback;
+import com.google.firebase.ai.type.ResponseModality;
 import com.google.firebase.ai.type.SafetyRating;
+import com.google.firebase.ai.type.SpeechConfig;
 import com.google.firebase.ai.type.TextPart;
 import com.google.firebase.ai.type.UsageMetadata;
-import com.google.firebase.concurrent.FirebaseExecutors;
+import com.google.firebase.ai.type.Voices;
 import java.util.Calendar;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
 import kotlinx.serialization.json.JsonElement;
 import kotlinx.serialization.json.JsonNull;
+import kotlinx.serialization.json.JsonObject;
 import org.junit.Assert;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
@@ -63,9 +75,31 @@ public class JavaCompileTests {
 
   public void initializeJava() throws Exception {
     FirebaseAI vertex = FirebaseAI.getInstance();
-    GenerativeModel model = vertex.generativeModel("fake-model-name");
+    GenerativeModel model = vertex.generativeModel("fake-model-name", getConfig());
+    LiveGenerativeModel live = vertex.liveModel("fake-model-name", getLiveConfig());
     GenerativeModelFutures futures = GenerativeModelFutures.from(model);
+    LiveModelFutures liveFutures = LiveModelFutures.from(live);
     testFutures(futures);
+    testLiveFutures(liveFutures);
+  }
+
+  private GenerationConfig getConfig() {
+    return new GenerationConfig.Builder().build();
+    // TODO b/406558430 GenerationConfig.Builder.setParts returns void
+  }
+
+  private LiveGenerationConfig getLiveConfig() {
+    return new LiveGenerationConfig.Builder()
+        .setTopK(10)
+        .setTopP(11.0F)
+        .setTemperature(32.0F)
+        .setCandidateCount(1)
+        .setMaxOutputTokens(0xCAFEBABE)
+        .setFrequencyPenalty(1.0F)
+        .setPresencePenalty(2.0F)
+        .setResponseModality(ResponseModality.AUDIO)
+        .setSpeechConfig(new SpeechConfig(Voices.AOEDE))
+        .build();
   }
 
   private void testFutures(GenerativeModelFutures futures) throws Exception {
@@ -138,7 +172,7 @@ public void validateGenerateContentResponse(GenerateContentResponse response) {
     List<Candidate> candidates = response.getCandidates();
     if (candidates.size() == 1
         && candidates.get(0).getContent().getParts().stream()
-            .anyMatch(p -> p instanceof TextPart && !((TextPart) p).getText().isEmpty())) {
+        .anyMatch(p -> p instanceof TextPart && !((TextPart) p).getText().isEmpty())) {
       String text = response.getText();
       Assert.assertNotNull(text);
       Assert.assertFalse(text.isBlank());
@@ -236,4 +270,62 @@ public void validateUsageMetadata(UsageMetadata metadata) {
       }
     }
   }
+
+  private void testLiveFutures(LiveModelFutures futures) throws Exception {
+    LiveSessionFutures session = futures.connect().get();
+    session
+        .receive()
+        .subscribe(
+            new Subscriber<LiveContentResponse>() {
+              @Override
+              public void onSubscribe(Subscription s) {
+                s.request(Long.MAX_VALUE);
+              }
+
+              @Override
+              public void onNext(LiveContentResponse response) {
+                validateLiveContentResponse(response);
+              }
+
+              @Override
+              public void onError(Throwable t) {
+                // Ignore
+              }
+
+              @Override
+              public void onComplete() {
+                // Also ignore
+              }
+            });
+
+    session.send("Fake message");
+    session.send(new Content.Builder().addText("Fake message").build());
+
+    byte[] bytes = new byte[] {(byte) 0xCA, (byte) 0xFE, (byte) 0xBA, (byte) 0xBE};
+    session.sendMediaStream(List.of(new MediaData(bytes, "image/jxl")));
+
+    FunctionResponsePart functionResponse =
+        new FunctionResponsePart("myFunction", new JsonObject(Map.of()));
+    session.sendFunctionResponse(List.of(functionResponse, functionResponse));
+
+    session.startAudioConversation(part -> functionResponse);
+    session.startAudioConversation();
+    session.stopAudioConversation();
+    session.stopReceiving();
+    session.close();
+  }
+
+  private void validateLiveContentResponse(LiveContentResponse response) {
+    //int status = response.getStatus();
+    //Assert.assertEquals(status, LiveContentResponse.Status.Companion.getNORMAL());
+    //Assert.assertNotEquals(status, LiveContentResponse.Status.Companion.getINTERRUPTED());
+    //Assert.assertNotEquals(status, LiveContentResponse.Status.Companion.getTURN_COMPLETE());
+    // TODO b/412743328 LiveContentResponse.Status inaccessible for Java users
+    Content data = response.getData();
+    if (data != null) {
+      validateContent(data);
+    }
+    String text = response.getText();
+    validateFunctionCalls(response.getFunctionCalls());
+  }
 }