feat: Support response MIME type in ChatFirebaseVertexAI (#461) (#463)

davidmigloz · Jun 15, 2024 · c345272 · c345272
1 parent c8b30c9
commit c345272
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 4 deletions.
diff --git a/packages/langchain_firebase/lib/src/chat_models/vertex_ai/chat_firebase_vertex_ai.dart b/packages/langchain_firebase/lib/src/chat_models/vertex_ai/chat_firebase_vertex_ai.dart
@@ -263,6 +263,12 @@ class ChatFirebaseVertexAI extends BaseChatModel<ChatFirebaseVertexAIOptions> {
         temperature: options?.temperature ?? defaultOptions.temperature,
         topP: options?.topP ?? defaultOptions.topP,
         topK: options?.topK ?? defaultOptions.topK,
+        responseMimeType:
+            options?.responseMimeType ?? defaultOptions.responseMimeType,
+        // responseSchema not supported yet
+        // responseSchema:
+        // (options?.responseSchema ?? defaultOptions.responseSchema)
+        //     ?.toSchema(),
       ),
       (options?.tools ?? defaultOptions.tools)?.toToolList(),
       (options?.toolChoice ?? defaultOptions.toolChoice)?.toToolConfig(),

diff --git a/packages/langchain_firebase/lib/src/chat_models/vertex_ai/mappers.dart b/packages/langchain_firebase/lib/src/chat_models/vertex_ai/mappers.dart
@@ -197,14 +197,17 @@ extension ChatToolListMapper on List<ToolSpec> {
           (tool) => f.FunctionDeclaration(
             tool.name,
             tool.description,
-            _mapJsonSchemaToSchema(tool.inputJsonSchema),
+            tool.inputJsonSchema.toSchema(),
           ),
         ).toList(growable: false),
       ),
     ];
   }
+}
 
-  f.Schema _mapJsonSchemaToSchema(final Map<String, dynamic> jsonSchema) {
+extension SchemaMapper on Map<String, dynamic> {
+  f.Schema toSchema() {
+    final jsonSchema = this;
     final type = jsonSchema['type'] as String;
     final description = jsonSchema['description'] as String?;
     final nullable = jsonSchema['nullable'] as bool?;
@@ -247,7 +250,7 @@ extension ChatToolListMapper on List<ToolSpec> {
         );
       case 'array':
         if (items != null) {
-          final itemsSchema = _mapJsonSchemaToSchema(items);
+          final itemsSchema = items.toSchema();
           return f.Schema.array(
             description: description,
             nullable: nullable,
@@ -258,7 +261,10 @@ extension ChatToolListMapper on List<ToolSpec> {
       case 'object':
         if (properties != null) {
           final propertiesSchema = properties.map(
-            (key, value) => MapEntry(key, _mapJsonSchemaToSchema(value)),
+            (key, value) => MapEntry(
+              key,
+              (value as Map<String, dynamic>).toSchema(),
+            ),
           );
           return f.Schema.object(
             properties: propertiesSchema,

diff --git a/packages/langchain_firebase/lib/src/chat_models/vertex_ai/types.dart b/packages/langchain_firebase/lib/src/chat_models/vertex_ai/types.dart
@@ -13,6 +13,7 @@ class ChatFirebaseVertexAIOptions extends ChatModelOptions {
     this.maxOutputTokens,
     this.temperature,
     this.stopSequences,
+    this.responseMimeType,
     this.safetySettings,
     super.tools,
     super.toolChoice,
@@ -69,6 +70,13 @@ class ChatFirebaseVertexAIOptions extends ChatModelOptions {
   /// The stop sequence will not be included as part of the response.
   final List<String>? stopSequences;
 
+  /// Output response mimetype of the generated candidate text.
+  ///
+  /// Supported mimetype:
+  /// - `text/plain`: (default) Text output.
+  /// - `application/json`: JSON response in the candidates.
+  final String? responseMimeType;
+
   /// A list of unique [ChatFirebaseVertexAISafetySetting] instances for blocking
   /// unsafe content.
   ///