refactor(java): expose chat and generativeModel

google-gemini · Dec 13, 2023 · e0d2499 · e0d2499
1 parent e891e0b
commit e0d2499
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
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
@@ -46,18 +46,18 @@ abstract class ChatFutures internal constructor() {
    */
   abstract fun sendMessageStream(prompt: Content): Publisher<GenerateContentResponse>
 
-  /** Returns the previous interactions with the model */
-  abstract fun getHistory(): ArrayList<Content>
+  /** Returns the [Chat] instance that was used to create this instance */
+  abstract fun getChat(): Chat
 
-  private class FuturesImpl(val chat: Chat) : ChatFutures() {
+  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 getHistory(): ArrayList<Content> {
-      return ArrayList(chat.history)
+    override fun getChat(): Chat {
+      return chat
     }
   }
 

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
@@ -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,10 @@ 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 {
+      return model
+    }
   }
 
   companion object {