LLM extraction allows top-level to be non-object (e.g. list) for OpenAI. (#302)

Author: badmonster0

examples/docs_to_kg/main.py

@@ -13,11 +13,6 @@ class Relationship:
     predicate: str
     object: str
 
-@dataclasses.dataclass
-class Relationships:
-    """Describe a relationship between two nodes."""
-    relationships: list[Relationship]
-
 @cocoindex.flow_def(name="DocsToKG")
 def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
     """
@@ -48,13 +43,16 @@ def docs_to_kg_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.D
                 cocoindex.functions.ExtractByLlm(
                     llm_spec=cocoindex.LlmSpec(
                         api_type=cocoindex.LlmApiType.OPENAI, model="gpt-4o"),
-                    output_type=Relationships,
+                    # Replace by this spec below, to use Ollama API instead of OpenAI
+                    #   llm_spec=cocoindex.LlmSpec(
+                    #       api_type=cocoindex.LlmApiType.OLLAMA, model="llama3.2"),
+                    output_type=list[Relationship],
                     instruction=(
                         "Please extract relationships from CocoIndex documents. "
                         "Focus on concepts and ingnore specific examples. "
                         "Each relationship should be a tuple of (subject, predicate, object).")))
 
-            with chunk["relationships"]["relationships"].row() as relationship:
+            with chunk["relationships"].row() as relationship:
                 relationship["subject_embedding"] = relationship["subject"].transform(
                     cocoindex.functions.SentenceTransformerEmbed(
                         model="sentence-transformers/all-MiniLM-L6-v2"))

src/base/json_schema.rs

@@ -1,8 +1,6 @@
-use crate::utils::immutable::RefList;
+use crate::prelude::*;
 
-use super::{schema, spec::FieldName};
-use anyhow::Result;
-use indexmap::IndexMap;
+use crate::utils::immutable::RefList;
 use schemars::schema::{
     ArrayValidation, InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec,
 };
@@ -19,6 +17,9 @@ pub struct ToJsonSchemaOptions {
 
     /// If true, extract descriptions to a separate extra instruction.
     pub extract_descriptions: bool,
+
+    /// If true, the top level must be a JSON object.
+    pub top_level_must_be_object: bool,
 }
 
 struct JsonSchemaBuilder {
@@ -38,7 +39,7 @@ impl JsonSchemaBuilder {
         &mut self,
         schema: &mut SchemaObject,
         description: impl ToString,
-        field_path: RefList<'_, &'_ FieldName>,
+        field_path: RefList<'_, &'_ spec::FieldName>,
     ) {
         if self.options.extract_descriptions {
             let mut fields: Vec<_> = field_path.iter().map(|f| f.as_str()).collect();
@@ -53,7 +54,7 @@ impl JsonSchemaBuilder {
     fn for_basic_value_type(
         &mut self,
         basic_type: &schema::BasicValueType,
-        field_path: RefList<'_, &'_ FieldName>,
+        field_path: RefList<'_, &'_ spec::FieldName>,
     ) -> SchemaObject {
         let mut schema = SchemaObject::default();
         match basic_type {
@@ -171,7 +172,7 @@ impl JsonSchemaBuilder {
     fn for_struct_schema(
         &mut self,
         struct_schema: &schema::StructSchema,
-        field_path: RefList<'_, &'_ FieldName>,
+        field_path: RefList<'_, &'_ spec::FieldName>,
     ) -> SchemaObject {
         let mut schema = SchemaObject::default();
         if let Some(description) = &struct_schema.description {
@@ -213,7 +214,7 @@ impl JsonSchemaBuilder {
     fn for_value_type(
         &mut self,
         value_type: &schema::ValueType,
-        field_path: RefList<'_, &'_ FieldName>,
+        field_path: RefList<'_, &'_ spec::FieldName>,
     ) -> SchemaObject {
         match value_type {
             schema::ValueType::Basic(b) => self.for_basic_value_type(b, field_path),
@@ -234,7 +235,7 @@ impl JsonSchemaBuilder {
     fn for_enriched_value_type(
         &mut self,
         enriched_value_type: &schema::EnrichedValueType,
-        field_path: RefList<'_, &'_ FieldName>,
+        field_path: RefList<'_, &'_ spec::FieldName>,
     ) -> SchemaObject {
         self.for_value_type(&enriched_value_type.typ, field_path)
     }
@@ -262,11 +263,69 @@ impl JsonSchemaBuilder {
     }
 }
 
+pub struct ValueExtractor {
+    value_type: schema::ValueType,
+    object_wrapper_field_name: Option<String>,
+}
+
+impl ValueExtractor {
+    pub fn extract_value(&self, json_value: serde_json::Value) -> Result<value::Value> {
+        let unwrapped_json_value =
+            if let Some(object_wrapper_field_name) = &self.object_wrapper_field_name {
+                match json_value {
+                    serde_json::Value::Object(mut o) => o
+                        .remove(object_wrapper_field_name)
+                        .unwrap_or(serde_json::Value::Null),
+                    _ => {
+                        bail!("Field `{}` not found", object_wrapper_field_name)
+                    }
+                }
+            } else {
+                json_value
+            };
+        let result = value::Value::from_json(unwrapped_json_value, &self.value_type)?;
+        Ok(result)
+    }
+}
+
+pub struct BuildJsonSchemaOutput {
+    pub schema: SchemaObject,
+    pub extra_instructions: Option<String>,
+    pub value_extractor: ValueExtractor,
+}
+
 pub fn build_json_schema(
-    value_type: &schema::EnrichedValueType,
+    value_type: schema::EnrichedValueType,
     options: ToJsonSchemaOptions,
-) -> Result<(SchemaObject, Option<String>)> {
+) -> Result<BuildJsonSchemaOutput> {
     let mut builder = JsonSchemaBuilder::new(options);
-    let schema = builder.for_enriched_value_type(value_type, RefList::Nil);
-    Ok((schema, builder.build_extra_instructions()?))
+    let (schema, object_wrapper_field_name) = if builder.options.top_level_must_be_object
+        && !matches!(value_type.typ, schema::ValueType::Struct(_))
+    {
+        let object_wrapper_field_name = "value".to_string();
+        let wrapper_struct = schema::StructSchema {
+            fields: Arc::new(vec![schema::FieldSchema {
+                name: object_wrapper_field_name.clone(),
+                value_type: value_type.clone(),
+            }]),
+            description: None,
+        };
+        (
+            builder.for_struct_schema(&wrapper_struct, RefList::Nil),
+            Some(object_wrapper_field_name),
+        )
+    } else {
+        (
+            builder.for_enriched_value_type(&value_type, RefList::Nil),
+            None,
+        )
+    };
+    Ok(BuildJsonSchemaOutput {
+        schema,
+        extra_instructions: builder.build_extra_instructions()?,
+        value_extractor: ValueExtractor {
+            value_type: value_type.typ,
+            object_wrapper_field_name,
+        },
+    })
 }

src/llm/ollama.rs

@@ -81,6 +81,7 @@ impl LlmGenerationClient for Client {
             fields_always_required: false,
             supports_format: true,
             extract_descriptions: true,
+            top_level_must_be_object: false,
         }
     }
 }

src/llm/openai.rs

@@ -103,6 +103,7 @@ impl LlmGenerationClient for Client {
             fields_always_required: true,
             supports_format: false,
             extract_descriptions: false,
+            top_level_must_be_object: true,
         }
     }
 }

src/ops/functions/extract_by_llm.rs

@@ -1,14 +1,12 @@
-use std::borrow::Cow;
-use std::sync::Arc;
-
-use schemars::schema::SchemaObject;
-use serde::Serialize;
+use crate::prelude::*;
 
-use crate::base::json_schema::build_json_schema;
 use crate::llm::{
     new_llm_generation_client, LlmGenerateRequest, LlmGenerationClient, LlmSpec, OutputFormat,
 };
 use crate::ops::sdk::*;
+use base::json_schema::build_json_schema;
+use schemars::schema::SchemaObject;
+use std::borrow::Cow;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Spec {
@@ -25,8 +23,8 @@ struct Executor {
     args: Args,
     client: Box<dyn LlmGenerationClient>,
     output_json_schema: SchemaObject,
-    output_type: EnrichedValueType,
     system_prompt: String,
+    value_extractor: base::json_schema::ValueExtractor,
 }
 
 fn get_system_prompt(instructions: &Option<String>, extra_instructions: Option<String>) -> String {
@@ -53,14 +51,13 @@ Output only the JSON without any additional messages or explanations."
 impl Executor {
     async fn new(spec: Spec, args: Args) -> Result<Self> {
         let client = new_llm_generation_client(spec.llm_spec).await?;
-        let (output_json_schema, extra_instructions) =
-            build_json_schema(&spec.output_type, client.json_schema_options())?;
+        let schema_output = build_json_schema(spec.output_type, client.json_schema_options())?;
         Ok(Self {
             args,
             client,
-            output_json_schema,
-            output_type: spec.output_type,
-            system_prompt: get_system_prompt(&spec.instruction, extra_instructions),
+            output_json_schema: schema_output.schema,
+            system_prompt: get_system_prompt(&spec.instruction, schema_output.extra_instructions),
+            value_extractor: schema_output.value_extractor,
         })
     }
 }
@@ -87,7 +84,7 @@ impl SimpleFunctionExecutor for Executor {
         };
         let res = self.client.generate(req).await?;
         let json_value: serde_json::Value = serde_json::from_str(res.text.as_str())?;
-        let value = Value::from_json(json_value, &self.output_type.typ)?;
+        let value = self.value_extractor.extract_value(json_value)?;
         Ok(value)
     }
 }

src/prelude.rs

@@ -17,7 +17,7 @@ pub(crate) use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 pub(crate) use std::hash::Hash;
 pub(crate) use std::sync::{Arc, LazyLock, Mutex, OnceLock, RwLock, Weak};
 
-pub(crate) use crate::base::{schema, spec, value};
+pub(crate) use crate::base::{self, schema, spec, value};
 pub(crate) use crate::builder::{self, plan};
 pub(crate) use crate::execution;
 pub(crate) use crate::lib_context::{get_lib_context, get_runtime, FlowContext, LibContext};