Structured JSON support

Cargo.toml

@@ -23,6 +23,7 @@ text-splitter = { version = "0.13.1", optional = true }
 regex = { version = "1.9.3", optional = true }
 async-stream = "0.3.5"
 http = {version = "1.1.0", optional = true }
+schemars = "0.8.21"
 
 [features]
 default = ["reqwest/default-tls"]

examples/structured_output.rs

@@ -0,0 +1,55 @@
+use ollama_rs::{
+    generation::{
+        completion::request::GenerationRequest,
+        options::GenerationOptions,
+        parameters::{schema_for, FormatType, JsonSchema, JsonStructure},
+    },
+    Ollama,
+};
+use serde::Deserialize;
+
+#[derive(JsonSchema, Deserialize, Debug)]
+enum Temperature {
+    Warm,
+    Cold,
+}
+
+#[allow(dead_code)]
+#[derive(JsonSchema, Deserialize, Debug)]
+struct Output {
+    country: String,
+    capital: String,
+    languages: Vec<String>,
+    temperature: Temperature,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let ollama = Ollama::default();
+    let model = "llama3.1:latest".to_string();
+    let prompt = "Tell me about the country north of the USA".to_string();
+
+    let format = FormatType::StructuredJson(JsonStructure::new::<Output>());
+    let res = ollama
+        .generate(
+            GenerationRequest::new(model, prompt)
+                .format(format)
+                .options(GenerationOptions::default().temperature(0.0)),
+        )
+        .await?;
+
+    let resp: Output = serde_json::from_str(&res.response)?;
+
+    // Output {
+    //     country: "Canada",
+    //     capital: "Ottawa",
+    //     languages: [
+    //         "English",
+    //         "French",
+    //     ],
+    //     temperature: Cold,
+    // }
+    dbg!(resp);
+
+    Ok(())
+}

src/generation/chat/request.rs

@@ -1,11 +1,11 @@
-use serde::{Deserialize, Serialize};
+use serde::Serialize;
 
 use crate::generation::{options::GenerationOptions, parameters::FormatType};
 
 use super::ChatMessage;
 
 /// A chat message request to Ollama.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize)]
 pub struct ChatMessageRequest {
     #[serde(rename = "model")]
     pub model_name: String,
@@ -41,7 +41,7 @@ impl ChatMessageRequest {
         self
     }
 
-    // The format to return a response in. Currently the only accepted value is `json`
+    /// The format to return a response in.
     pub fn format(mut self, format: FormatType) -> Self {
         self.format = Some(format);
         self

src/generation/completion/request.rs

@@ -91,7 +91,7 @@ impl GenerationRequest {
         self
     }
 
-    // The format to return a response in. Currently the only accepted value is `json`
+    /// The format to return a response in.
     pub fn format(mut self, format: FormatType) -> Self {
         self.format = Some(format);
         self

src/generation/functions/request.rs

@@ -33,7 +33,7 @@ impl FunctionCallRequest {
         self
     }
 
-    // The format to return a response in. Currently the only accepted value is `json`
+    /// The format to return a response in.
     pub fn format(mut self, format: FormatType) -> Self {
         self.chat.format = Some(format);
         self

src/generation/parameters.rs

@@ -1,10 +1,50 @@
-use serde::{Deserialize, Serialize};
+use schemars::{gen::SchemaSettings, schema::RootSchema};
+pub use schemars::{schema_for, JsonSchema};
+use serde::{Serialize, Serializer};
 
-/// The format to return a response in. Currently the only accepted value is `json`
-#[derive(Debug, Serialize, Deserialize, Clone)]
-#[serde(rename_all = "lowercase")]
+/// The format to return a response in.
+#[derive(Debug, Clone)]
 pub enum FormatType {
     Json,
+
+    /// Requires Ollama 0.5.0 or greater.
+    StructuredJson(JsonStructure),
+}
+
+impl Serialize for FormatType {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match self {
+            FormatType::Json => serializer.serialize_str("json"),
+            FormatType::StructuredJson(s) => s.schema.serialize(serializer),
+        }
+    }
+}
+
+/// Represents a serialized JSON schema. You can create this by converting
+/// a JsonSchema:
+/// ```rust
+/// let json_schema = schema_for!(Output);
+/// let serialized: SerializedJsonSchema = json_schema.into();
+/// ```
+#[derive(Debug, Clone)]
+pub struct JsonStructure {
+    schema: RootSchema,
+}
+
+impl JsonStructure {
+    pub fn new<T: JsonSchema>() -> Self {
+        // Need to do this because Ollama doesn't support $refs (references in the schema)
+        // So we have to explicitly turn them off
+        let mut settings = SchemaSettings::draft07();
+        settings.inline_subschemas = true;
+        let generator = settings.into_generator();
+        let schema = generator.into_root_schema_for::<T>();
+
+        Self { schema }
+    }
 }
 
 /// Used to control how long a model stays loaded in memory, by default models are unloaded after 5 minutes of inactivity