From 440f960d8ddd96b846ded23d7f20a983b635fa03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ros=C3=A1rio=20P=2E=20Fernandes?= <rosariopf@google.com>
Date: Thu, 14 Dec 2023 21:36:06 +0000
Subject: [PATCH] feat(java): expose the instances used to create Futures (#5)

This adds a `ChatFutures.getChatHistory()` helper method for Java
developers to be able to obtain the chat history. This new method should
be equivalent to Kotlin's `Chat.history`.
---
 .changes/cemetery-caption-bird-authority.json              | 1 +
 .../com/google/ai/client/generativeai/java/ChatFutures.kt  | 7 ++++++-
 .../ai/client/generativeai/java/GenerativeModelFutures.kt  | 7 ++++++-
 3 files changed, 13 insertions(+), 2 deletions(-)
 create mode 100644 .changes/cemetery-caption-bird-authority.json

diff --git a/.changes/cemetery-caption-bird-authority.json b/.changes/cemetery-caption-bird-authority.json
new file mode 100644
index 00000000..e23972af
--- /dev/null
+++ b/.changes/cemetery-caption-bird-authority.json
@@ -0,0 +1 @@
+{"type":"MINOR","changes":["Expose the instances used to create Future APIs in java"]}
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/java/ChatFutures.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/java/ChatFutures.kt
index 615f9cbd..81669c3f 100644
--- a/generativeai/src/main/java/com/google/ai/client/generativeai/java/ChatFutures.kt
+++ b/generativeai/src/main/java/com/google/ai/client/generativeai/java/ChatFutures.kt
@@ -46,12 +46,17 @@ abstract class ChatFutures internal constructor() {
    */
   abstract fun sendMessageStream(prompt: Content): Publisher<GenerateContentResponse>
 
-  private class FuturesImpl(val chat: Chat) : ChatFutures() {
+  /** Returns the [Chat] instance that was used to create this instance */
+  abstract fun getChat(): Chat
+
+  private class FuturesImpl(private val chat: Chat) : ChatFutures() {
     override fun sendMessage(prompt: Content): ListenableFuture<GenerateContentResponse> =
       SuspendToFutureAdapter.launchFuture { chat.sendMessage(prompt) }
 
     override fun sendMessageStream(prompt: Content): Publisher<GenerateContentResponse> =
       chat.sendMessageStream(prompt).asPublisher()
+
+    override fun getChat(): Chat = chat
   }
 
   companion object {
diff --git a/generativeai/src/main/java/com/google/ai/client/generativeai/java/GenerativeModelFutures.kt b/generativeai/src/main/java/com/google/ai/client/generativeai/java/GenerativeModelFutures.kt
index 36d2b34a..e1902f0c 100644
--- a/generativeai/src/main/java/com/google/ai/client/generativeai/java/GenerativeModelFutures.kt
+++ b/generativeai/src/main/java/com/google/ai/client/generativeai/java/GenerativeModelFutures.kt
@@ -64,7 +64,10 @@ abstract class GenerativeModelFutures internal constructor() {
    */
   abstract fun startChat(history: List<Content>): ChatFutures
 
-  private class FuturesImpl(val model: GenerativeModel) : GenerativeModelFutures() {
+  /** Returns the [GenerativeModel] instance that was used to create this object */
+  abstract fun getGenerativeModel(): GenerativeModel
+
+  private class FuturesImpl(private val model: GenerativeModel) : GenerativeModelFutures() {
     override fun generateContent(
       vararg prompt: Content
     ): ListenableFuture<GenerateContentResponse> =
@@ -79,6 +82,8 @@ abstract class GenerativeModelFutures internal constructor() {
     override fun startChat(): ChatFutures = startChat(emptyList())
 
     override fun startChat(history: List<Content>): ChatFutures = from(model.startChat(history))
+
+    override fun getGenerativeModel(): GenerativeModel = model
   }
 
   companion object {