From dc1b5cabdf3e8d11db2be0a19ce584c39c75c6cf Mon Sep 17 00:00:00 2001
From: Andre Bandarra <andreban@gmail.com>
Date: Fri, 20 Sep 2024 15:47:21 +0100
Subject: [PATCH] Adds support for the response schema field

- Support is added by allowing a serde_json `Value`.
- In the future, this could be improved by using a proper `Schema`
definition.
---
 examples/json-schema.rs       | 52 +++++++++++++++++++++++++++++++++++
 src/types/generate_content.rs |  3 ++
 2 files changed, 55 insertions(+)
 create mode 100644 examples/json-schema.rs

diff --git a/examples/json-schema.rs b/examples/json-schema.rs
new file mode 100644
index 0000000..748633a
--- /dev/null
+++ b/examples/json-schema.rs
@@ -0,0 +1,52 @@
+use gemini_rs::prelude::*;
+use serde_json::json;
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let authentication_manager = gcp_auth::provider().await?;
+    let api_endpoint = std::env::var("API_ENDPOINT")?;
+    let project_id = std::env::var("PROJECT_ID")?;
+    let location_id = std::env::var("LOCATION_ID")?;
+
+    let gemini = GeminiClient::new(
+        authentication_manager,
+        api_endpoint,
+        project_id,
+        location_id,
+    );
+
+    let prompt = "Generate 10 ideas of blog posts with a title and decription for each idea.";
+    let request = GenerateContentRequest {
+        contents: vec![Content {
+            role: Some("user".to_string()),
+            parts: Some(vec![Part::Text(prompt.to_string())]),
+        }],
+        generation_config: Some(GenerationConfig {
+            response_mime_type: Some("application/json".to_string()),
+            response_schema: Some(json!({
+                "type": "ARRAY",
+                "items": {
+                    "type": "OBJECT",
+                    "properties": {
+                        "title": {
+                            "type": "STRING"
+                        },
+                        "description": {
+                            "type": "STRING"
+                        }
+                    }
+                }
+            })),
+            ..Default::default()
+        }),
+        ..Default::default()
+    };
+
+    let result = gemini
+        .generate_content(&request, "gemini-1.5-flash-001")
+        .await?;
+
+    println!("Response: {:?}", result.candidates[0].get_text().unwrap());
+
+    Ok(())
+}
diff --git a/src/types/generate_content.rs b/src/types/generate_content.rs
index 8dd1194..99c2df8 100644
--- a/src/types/generate_content.rs
+++ b/src/types/generate_content.rs
@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 
 use serde::{Deserialize, Serialize};
+use serde_json::Value;
 
 use super::{Content, Part, VertexApiError};
 use crate::error::Result;
@@ -66,6 +67,8 @@ pub struct GenerationConfig {
     pub candidate_count: Option<u32>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub response_mime_type: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub response_schema: Option<Value>,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]