feat: Support response MIME type and schema in ChatGoogleGenerativeAI (…

…#461)
davidmigloz · Jun 15, 2024 · e258399 · e258399
1 parent 6d13729
commit e258399
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 6 deletions.
diff --git a/melos.yaml b/melos.yaml
@@ -41,7 +41,7 @@ command:
       flutter_markdown: ^0.6.22
       freezed_annotation: ^2.4.1
       gcloud: ^0.8.12
-      google_generative_ai: 0.4.0
+      google_generative_ai: 0.4.3
       googleapis: ^12.0.0
       googleapis_auth: ^1.5.1
       http: ^1.1.0

diff --git a/packages/langchain_google/lib/src/chat_models/google_ai/chat_google_generative_ai.dart b/packages/langchain_google/lib/src/chat_models/google_ai/chat_google_generative_ai.dart
@@ -322,6 +322,11 @@ class ChatGoogleGenerativeAI
         temperature: options?.temperature ?? defaultOptions.temperature,
         topP: options?.topP ?? defaultOptions.topP,
         topK: options?.topK ?? defaultOptions.topK,
+        responseMimeType:
+            options?.responseMimeType ?? defaultOptions.responseMimeType,
+        responseSchema:
+            (options?.responseSchema ?? defaultOptions.responseSchema)
+                ?.toSchema(),
       ),
       (options?.tools ?? defaultOptions.tools)?.toToolList(),
       (options?.toolChoice ?? defaultOptions.toolChoice)?.toToolConfig(),

diff --git a/packages/langchain_google/lib/src/chat_models/google_ai/mappers.dart b/packages/langchain_google/lib/src/chat_models/google_ai/mappers.dart
@@ -198,14 +198,17 @@ extension ChatToolListMapper on List<ToolSpec> {
           (tool) => g.FunctionDeclaration(
             tool.name,
             tool.description,
-            _mapJsonSchemaToSchema(tool.inputJsonSchema),
+            tool.inputJsonSchema.toSchema(),
           ),
         ).toList(growable: false),
       ),
     ];
   }
+}
 
-  g.Schema _mapJsonSchemaToSchema(final Map<String, dynamic> jsonSchema) {
+extension SchemaMapper on Map<String, dynamic> {
+  g.Schema toSchema() {
+    final jsonSchema = this;
     final type = jsonSchema['type'] as String;
     final description = jsonSchema['description'] as String?;
     final nullable = jsonSchema['nullable'] as bool?;
@@ -248,7 +251,7 @@ extension ChatToolListMapper on List<ToolSpec> {
         );
       case 'array':
         if (items != null) {
-          final itemsSchema = _mapJsonSchemaToSchema(items);
+          final itemsSchema = items.toSchema();
           return g.Schema.array(
             items: itemsSchema,
             description: description,
@@ -259,7 +262,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 g.Schema.object(
             properties: propertiesSchema,

diff --git a/packages/langchain_google/lib/src/chat_models/google_ai/types.dart b/packages/langchain_google/lib/src/chat_models/google_ai/types.dart
@@ -13,6 +13,8 @@ class ChatGoogleGenerativeAIOptions extends ChatModelOptions {
     this.maxOutputTokens,
     this.temperature,
     this.stopSequences,
+    this.responseMimeType,
+    this.responseSchema,
     this.safetySettings,
     super.tools,
     super.toolChoice,
@@ -68,6 +70,39 @@ class ChatGoogleGenerativeAIOptions 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;
+
+  /// Output response schema of the generated candidate text.
+  /// Following the [JSON Schema specification](https://json-schema.org).
+  ///
+  /// - Note: This only applies when the specified ``responseMIMEType`` supports
+  ///   a schema; currently this is limited to `application/json`.
+  ///
+  /// Example:
+  /// ```json
+  /// {
+  ///   'type': 'object',
+  ///   'properties': {
+  ///     'answer': {
+  ///       'type': 'string',
+  ///       'description': 'The answer to the question being asked',
+  ///     },
+  ///     'sources': {
+  ///       'type': 'array',
+  ///       'items': {'type': 'string'},
+  ///       'description': 'The sources used to answer the question',
+  ///     },
+  ///   },
+  ///   'required': ['answer', 'sources'],
+  /// },
+  /// ```
+  final Map<String, dynamic>? responseSchema;
+
   /// A list of unique [ChatGoogleGenerativeAISafetySetting] instances for blocking
   /// unsafe content.
   ///

diff --git a/packages/langchain_google/pubspec.yaml b/packages/langchain_google/pubspec.yaml
@@ -20,7 +20,7 @@ dependencies:
   collection: '>=1.17.0 <1.19.0'
   fetch_client: ^1.0.2
   gcloud: ^0.8.12
-  google_generative_ai: 0.4.0
+  google_generative_ai: 0.4.3
   googleapis: ^12.0.0
   googleapis_auth: ^1.5.1
   http: ^1.1.0