diff --git a/firestore-genai-chatbot/extension.yaml b/firestore-genai-chatbot/extension.yaml
index 99162420..9937cf51 100644
--- a/firestore-genai-chatbot/extension.yaml
+++ b/firestore-genai-chatbot/extension.yaml
@@ -258,6 +258,22 @@ params:
     required: false
     immutable: false
 
+  - param: RESPONSE_MIME_TYPE
+    label: Output response mimetype of the generated candidate text.
+    description: >-
+      Supported mimetype:
+      - text/plain: (default) Text output.
+      - application/json: JSON response in the candidates.
+      The model needs to be prompted to output the appropriate response type, otherwise the behavior is undefined.
+    type: select
+    options:
+      - label: JSON
+        value: application/json
+      - label: Text
+        value: text/plain
+    default: text/plain
+    required: false
+
   - param: CANDIDATE_COUNT
     label: Candidate count
     description: >-
diff --git a/firestore-genai-chatbot/functions/package.json b/firestore-genai-chatbot/functions/package.json
index d937febb..1783242e 100644
--- a/firestore-genai-chatbot/functions/package.json
+++ b/firestore-genai-chatbot/functions/package.json
@@ -14,7 +14,7 @@
     "@google-ai/generativelanguage": "^2.3.0",
     "@google-cloud/aiplatform": "^3.19.0",
     "@google-cloud/vertexai": "^1.1.0",
-    "@google/generative-ai": "^0.1.1",
+    "@google/generative-ai": "^0.11.4",
     "@types/jest": "^29.5.12",
     "firebase-admin": "^12.1.0",
     "firebase-functions": "^4.9.0",
diff --git a/firestore-genai-chatbot/functions/src/config.ts b/firestore-genai-chatbot/functions/src/config.ts
index 969d2ca0..e1998239 100644
--- a/firestore-genai-chatbot/functions/src/config.ts
+++ b/firestore-genai-chatbot/functions/src/config.ts
@@ -48,6 +48,7 @@ export interface Config {
   safetySettings?: GoogleAISafetySetting[] | VertexSafetySetting[];
   provider: GenerativeAIProvider;
   maxOutputTokens?: number;
+  responseMimeType?: string;
 }
 
 function getSafetySettings(): GoogleAISafetySetting[] | VertexSafetySetting[] {
@@ -116,6 +117,7 @@ const config: Config = {
   maxOutputTokens: process.env.MAX_OUTPUT_TOKENS
     ? parseInt(process.env.MAX_OUTPUT_TOKENS)
     : undefined,
+  responseMimeType: process.env.RESPONSE_MIME_TYPE || 'text/plain',
 };
 
 export default config;
diff --git a/firestore-genai-chatbot/functions/src/generative-client/google_ai.ts b/firestore-genai-chatbot/functions/src/generative-client/google_ai.ts
index 44286526..e8e9d91d 100644
--- a/firestore-genai-chatbot/functions/src/generative-client/google_ai.ts
+++ b/firestore-genai-chatbot/functions/src/generative-client/google_ai.ts
@@ -1,5 +1,5 @@
 import {DiscussionClient, Message} from './base_class';
-import {GoogleGenerativeAI, InputContent} from '@google/generative-ai';
+import {GoogleGenerativeAI, Content} from '@google/generative-ai';
 import {logger} from 'firebase-functions/v1';
 import {SafetySetting} from '@google/generative-ai';
 
@@ -15,6 +15,7 @@ interface GeminiChatOptions {
   location: string;
   context?: string;
   safetySettings: SafetySetting[];
+  responseMimeType?: string;
 }
 
 type ApiMessage = {
@@ -80,6 +81,7 @@ export class GeminiDiscussionClient extends DiscussionClient<
         temperature: options.temperature,
         maxOutputTokens: options.maxOutputTokens,
         candidateCount: options.candidateCount,
+        responseMimeType: options.responseMimeType,
       },
       safetySettings: options.safetySettings,
     });
@@ -112,7 +114,7 @@ export class GeminiDiscussionClient extends DiscussionClient<
   }
 
   private messagesToApi(messages: Message[]) {
-    const out: InputContent[] = [];
+    const out: Content[] = [];
     for (const message of messages) {
       if (!message.prompt || !message.response) {
         // logs.warnMissingPromptOrResponse(message.path!);