Dump description as extra output for ollama.

Author: badmonster0

src/base/json_schema.rs

@@ -1,7 +1,12 @@
-use super::schema;
+use crate::utils::immutable::RefList;
+
+use super::{schema, spec::FieldName};
+use anyhow::Result;
+use indexmap::IndexMap;
 use schemars::schema::{
-    ArrayValidation, InstanceType, Metadata, ObjectValidation, Schema, SchemaObject, SingleOrVec,
+    ArrayValidation, InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec,
 };
+use std::fmt::Write;
 
 pub struct ToJsonSchemaOptions {
     /// If true, mark all fields as required.
@@ -11,18 +16,45 @@ pub struct ToJsonSchemaOptions {
 
     /// If true, the JSON schema supports the `format` keyword.
     pub supports_format: bool,
+
+    /// If true, extract descriptions to a separate extra instruction.
+    pub extract_descriptions: bool,
 }
 
-pub struct JsonSchemaBuilder {
+struct JsonSchemaBuilder {
     options: ToJsonSchemaOptions,
+    extra_instructions_per_field: IndexMap<String, String>,
 }
 
 impl JsonSchemaBuilder {
-    pub fn new(options: ToJsonSchemaOptions) -> Self {
-        Self { options }
+    fn new(options: ToJsonSchemaOptions) -> Self {
+        Self {
+            options,
+            extra_instructions_per_field: IndexMap::new(),
+        }
+    }
+
+    fn set_description(
+        &mut self,
+        schema: &mut SchemaObject,
+        description: impl ToString,
+        field_path: RefList<'_, &'_ FieldName>,
+    ) {
+        if self.options.extract_descriptions {
+            let mut fields: Vec<_> = field_path.iter().map(|f| f.as_str()).collect();
+            fields.reverse();
+            self.extra_instructions_per_field
+                .insert(fields.join("."), description.to_string());
+        } else {
+            schema.metadata.get_or_insert_default().description = Some(description.to_string());
+        }
     }
 
-    fn for_basic_value_type(&mut self, basic_type: &schema::BasicValueType) -> SchemaObject {
+    fn for_basic_value_type(
+        &mut self,
+        basic_type: &schema::BasicValueType,
+        field_path: RefList<'_, &'_ FieldName>,
+    ) -> SchemaObject {
         let mut schema = SchemaObject::default();
         match basic_type {
             schema::BasicValueType::Str => {
@@ -56,48 +88,66 @@ impl JsonSchemaBuilder {
                     max_items: Some(2),
                     ..Default::default()
                 }));
-                schema.metadata.get_or_insert_default().description =
-                    Some("A range, start pos (inclusive), end pos (exclusive).".to_string());
+                self.set_description(
+                    &mut schema,
+                    "A range represented by a list of two positions, start pos (inclusive), end pos (exclusive).",
+                    field_path,
+                );
             }
             schema::BasicValueType::Uuid => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
                 if self.options.supports_format {
                     schema.format = Some("uuid".to_string());
                 }
-                schema.metadata.get_or_insert_default().description =
-                    Some("A UUID, e.g. 123e4567-e89b-12d3-a456-426614174000".to_string());
+                self.set_description(
+                    &mut schema,
+                    "A UUID, e.g. 123e4567-e89b-12d3-a456-426614174000",
+                    field_path,
+                );
             }
             schema::BasicValueType::Date => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
                 if self.options.supports_format {
                     schema.format = Some("date".to_string());
                 }
-                schema.metadata.get_or_insert_default().title =
-                    Some("A date in YYYY-MM-DD format, e.g. 2025-03-27".to_string());
+                self.set_description(
+                    &mut schema,
+                    "A date in YYYY-MM-DD format, e.g. 2025-03-27",
+                    field_path,
+                );
             }
             schema::BasicValueType::Time => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
                 if self.options.supports_format {
                     schema.format = Some("time".to_string());
                 }
-                schema.metadata.get_or_insert_default().description =
-                    Some("A time in HH:MM:SS format, e.g. 13:32:12".to_string());
+                self.set_description(
+                    &mut schema,
+                    "A time in HH:MM:SS format, e.g. 13:32:12",
+                    field_path,
+                );
             }
             schema::BasicValueType::LocalDateTime => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
                 if self.options.supports_format {
                     schema.format = Some("date-time".to_string());
                 }
-                schema.metadata.get_or_insert_default().description =
-                    Some("Date time without timezone offset in YYYY-MM-DDTHH:MM:SS format, e.g. 2025-03-27T13:32:12".to_string());
+                self.set_description(
+                    &mut schema,
+                    "Date time without timezone offset in YYYY-MM-DDTHH:MM:SS format, e.g. 2025-03-27T13:32:12",
+                    field_path,
+                );
             }
             schema::BasicValueType::OffsetDateTime => {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::String)));
                 if self.options.supports_format {
                     schema.format = Some("date-time".to_string());
                 }
-                schema.metadata.get_or_insert_default().description =
-                    Some("Date time with timezone offset in RFC3339, e.g. 2025-03-27T13:32:12Z, 2025-03-27T07:32:12.313-06:00".to_string());
+                self.set_description(
+                    &mut schema,
+                    "Date time with timezone offset in RFC3339, e.g. 2025-03-27T13:32:12Z, 2025-03-27T07:32:12.313-06:00",
+                    field_path,
+                );
             }
             schema::BasicValueType::Json => {
                 // Can be any value. No type constraint.
@@ -106,7 +156,8 @@ impl JsonSchemaBuilder {
                 schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Array)));
                 schema.array = Some(Box::new(ArrayValidation {
                     items: Some(SingleOrVec::Single(Box::new(
-                        self.for_basic_value_type(&s.element_type).into(),
+                        self.for_basic_value_type(&s.element_type, field_path)
+                            .into(),
                     ))),
                     min_items: s.dimension.and_then(|d| u32::try_from(d).ok()),
                     max_items: s.dimension.and_then(|d| u32::try_from(d).ok()),
@@ -117,54 +168,61 @@ impl JsonSchemaBuilder {
         schema
     }
 
-    fn for_struct_schema(&mut self, struct_schema: &schema::StructSchema) -> SchemaObject {
-        SchemaObject {
-            metadata: Some(Box::new(Metadata {
-                description: struct_schema.description.as_ref().map(|s| s.to_string()),
-                ..Default::default()
-            })),
-            instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::Object))),
-            object: Some(Box::new(ObjectValidation {
-                properties: struct_schema
-                    .fields
-                    .iter()
-                    .map(|f| {
-                        let mut schema = self.for_enriched_value_type(&f.value_type);
-                        if self.options.fields_always_required && f.value_type.nullable {
-                            if let Some(instance_type) = &mut schema.instance_type {
-                                let mut types = match instance_type {
-                                    SingleOrVec::Single(t) => vec![**t],
-                                    SingleOrVec::Vec(t) => std::mem::take(t),
-                                };
-                                types.push(InstanceType::Null);
-                                *instance_type = SingleOrVec::Vec(types);
-                            }
+    fn for_struct_schema(
+        &mut self,
+        struct_schema: &schema::StructSchema,
+        field_path: RefList<'_, &'_ FieldName>,
+    ) -> SchemaObject {
+        let mut schema = SchemaObject::default();
+        if let Some(description) = &struct_schema.description {
+            self.set_description(&mut schema, description, field_path);
+        }
+        schema.instance_type = Some(SingleOrVec::Single(Box::new(InstanceType::Object)));
+        schema.object = Some(Box::new(ObjectValidation {
+            properties: struct_schema
+                .fields
+                .iter()
+                .map(|f| {
+                    let mut schema =
+                        self.for_enriched_value_type(&f.value_type, field_path.prepend(&f.name));
+                    if self.options.fields_always_required && f.value_type.nullable {
+                        if let Some(instance_type) = &mut schema.instance_type {
+                            let mut types = match instance_type {
+                                SingleOrVec::Single(t) => vec![**t],
+                                SingleOrVec::Vec(t) => std::mem::take(t),
+                            };
+                            types.push(InstanceType::Null);
+                            *instance_type = SingleOrVec::Vec(types);
                         }
-                        (f.name.to_string(), schema.into())
-                    })
-                    .collect(),
-                required: struct_schema
-                    .fields
-                    .iter()
-                    .filter(|&f| (self.options.fields_always_required || !f.value_type.nullable))
-                    .map(|f| f.name.to_string())
-                    .collect(),
-                additional_properties: Some(Schema::Bool(false).into()),
-                ..Default::default()
-            })),
+                    }
+                    (f.name.to_string(), schema.into())
+                })
+                .collect(),
+            required: struct_schema
+                .fields
+                .iter()
+                .filter(|&f| (self.options.fields_always_required || !f.value_type.nullable))
+                .map(|f| f.name.to_string())
+                .collect(),
+            additional_properties: Some(Schema::Bool(false).into()),
             ..Default::default()
-        }
+        }));
+        schema
     }
 
-    fn for_value_type(&mut self, value_type: &schema::ValueType) -> SchemaObject {
+    fn for_value_type(
+        &mut self,
+        value_type: &schema::ValueType,
+        field_path: RefList<'_, &'_ FieldName>,
+    ) -> SchemaObject {
         match value_type {
-            schema::ValueType::Basic(b) => self.for_basic_value_type(b),
-            schema::ValueType::Struct(s) => self.for_struct_schema(s),
+            schema::ValueType::Basic(b) => self.for_basic_value_type(b, field_path),
+            schema::ValueType::Struct(s) => self.for_struct_schema(s, field_path),
             schema::ValueType::Collection(c) => SchemaObject {
                 instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::Array))),
                 array: Some(Box::new(ArrayValidation {
                     items: Some(SingleOrVec::Single(Box::new(
-                        self.for_struct_schema(&c.row).into(),
+                        self.for_struct_schema(&c.row, field_path).into(),
                     ))),
                     ..Default::default()
                 })),
@@ -173,10 +231,42 @@ impl JsonSchemaBuilder {
         }
     }
 
-    pub fn for_enriched_value_type(
+    fn for_enriched_value_type(
         &mut self,
         enriched_value_type: &schema::EnrichedValueType,
+        field_path: RefList<'_, &'_ FieldName>,
     ) -> SchemaObject {
-        self.for_value_type(&enriched_value_type.typ)
+        self.for_value_type(&enriched_value_type.typ, field_path)
     }
+
+    fn build_extra_instructions(&self) -> Result<Option<String>> {
+        if self.extra_instructions_per_field.is_empty() {
+            return Ok(None);
+        }
+
+        let mut instructions = String::new();
+        write!(&mut instructions, "Instructions for specific fields:\n\n")?;
+        for (field_path, instruction) in self.extra_instructions_per_field.iter() {
+            write!(
+                &mut instructions,
+                "- {}: {}\n\n",
+                if field_path.is_empty() {
+                    "(root object)"
+                } else {
+                    field_path.as_str()
+                },
+                instruction
+            )?;
+        }
+        Ok(Some(instructions))
+    }
+}
+
+pub fn build_json_schema(
+    value_type: &schema::EnrichedValueType,
+    options: ToJsonSchemaOptions,
+) -> Result<(SchemaObject, Option<String>)> {
+    let mut builder = JsonSchemaBuilder::new(options);
+    let schema = builder.for_enriched_value_type(value_type, RefList::Nil);
+    Ok((schema, builder.build_extra_instructions()?))
 }

src/base/value.rs

@@ -1,7 +1,7 @@
 use crate::{api_bail, api_error};
 
 use super::schema::*;
-use anyhow::{Context, Result};
+use anyhow::Result;
 use base64::prelude::*;
 use chrono::Offset;
 use log::warn;
@@ -10,7 +10,7 @@ use serde::{
     ser::{SerializeMap, SerializeSeq, SerializeTuple},
     Deserialize, Serialize,
 };
-use std::{collections::BTreeMap, ops::Deref, str::FromStr, sync::Arc};
+use std::{collections::BTreeMap, ops::Deref, sync::Arc};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct RangeValue {

src/llm/mod.rs

@@ -47,24 +47,7 @@ pub trait LlmGenerationClient: Send + Sync {
         request: LlmGenerateRequest<'req>,
     ) -> Result<LlmGenerateResponse>;
 
-    /// If true, the LLM only accepts a JSON schema with all fields required.
-    /// This is a limitation of LLM models such as OpenAI.
-    /// Otherwise, the LLM will accept a JSON schema with optional fields.
-    fn json_schema_fields_always_required(&self) -> bool {
-        false
-    }
-
-    /// If true, the LLM supports the `format` keyword in the JSON schema.
-    fn json_schema_supports_format(&self) -> bool {
-        true
-    }
-
-    fn to_json_schema_options(&self) -> ToJsonSchemaOptions {
-        ToJsonSchemaOptions {
-            fields_always_required: self.json_schema_fields_always_required(),
-            supports_format: self.json_schema_supports_format(),
-        }
-    }
+    fn json_schema_options(&self) -> ToJsonSchemaOptions;
 }
 
 mod ollama;

src/llm/ollama.rs

@@ -75,4 +75,12 @@ impl LlmGenerationClient for Client {
             text: json.response,
         })
     }
+
+    fn json_schema_options(&self) -> super::ToJsonSchemaOptions {
+        super::ToJsonSchemaOptions {
+            fields_always_required: false,
+            supports_format: true,
+            extract_descriptions: true,
+        }
+    }
 }

src/llm/openai.rs

@@ -98,11 +98,11 @@ impl LlmGenerationClient for Client {
         Ok(super::LlmGenerateResponse { text })
     }
 
-    fn json_schema_fields_always_required(&self) -> bool {
-        true
-    }
-
-    fn json_schema_supports_format(&self) -> bool {
-        false
+    fn json_schema_options(&self) -> super::ToJsonSchemaOptions {
+        super::ToJsonSchemaOptions {
+            fields_always_required: true,
+            supports_format: false,
+            extract_descriptions: false,
+        }
     }
 }