Update app.py

app.py

@@ -2,9 +2,9 @@ import os
 import argparse
 import gradio as gr
 from difflib import Differ
-from functools import partial
 from string import Template
 from utils import load_prompt, setup_gemini_client
+from configs.responses import SummaryResponses
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -26,8 +26,8 @@ def find_attached_file(filename, attached_files):
     return None
 
 def echo(message, history, state):
-
     attached_file = None
+    
     if message['files']:
         path_local = message['files'][0]
         filename = os.path.basename(path_local)
@@ -45,8 +45,6 @@ def echo(message, history, state):
             })
             attached_file = path_gcp
 
-    # [{'role': 'user', 'metadata': None, 'content': 'asdf', 'options': None}, {'role': 'assistant', 'metadata': None, 'content': 'asdf', 'options': None}]
-
     user_message = [message['text']]
     if attached_file: user_message.append(attached_file)
 
@@ -56,7 +54,7 @@ def echo(message, history, state):
 
     response = client.models.generate_content(
         model="gemini-1.5-flash",
-        contents=state['messages']
+        contents=state['messages'],
     )
     model_response = response.text
 
@@ -71,7 +69,10 @@ def echo(message, history, state):
                     previous_summary=state['summary'], 
                     latest_conversation=str({"user": message['text'], "assistant": model_response})
                 )
-            ]
+            ],
+            config={'response_mime_type': 'application/json',
+                'response_schema': SummaryResponses,
+            },
         )
 
     if state['summary'] != "":
@@ -80,8 +81,16 @@ def echo(message, history, state):
         prev_summary = ""
 
     d = Differ()
-    state['summary'] = response.text
-    state['summary_history'].append(response.text)
+    state['summary'] = (
+        response.parsed.summary 
+        if getattr(response.parsed, "summary", None) is not None 
+        else response.text
+    )
+    state['summary_history'].append(
+        response.parsed.summary 
+        if getattr(response.parsed, "summary", None) is not None 
+        else response.text
+    )
     state['summary_diff_history'].append(
         [
             (token[2:], token[0] if token[0] != " " else None)
@@ -188,4 +197,4 @@ def main(args):
 if __name__ == "__main__":
     args = parse_args()
     demo = main(args)
-    demo.launch()
+    demo.launch()