Repository: tmc/langchaingo
Pull request #270 (2023-08-29 02:37:34)

Add jsonschema
This pulls in jsonschema for now, until we incorporate #145


Patch:
diff --git a/jsonschema/json.go b/jsonschema/json.go
new file mode 100644
index 00000000..cb941eb7
--- /dev/null
+++ b/jsonschema/json.go
@@ -0,0 +1,49 @@
+// Package jsonschema provides very simple functionality for representing a JSON schema as a
+// (nested) struct. This struct can be used with the chat completion "function call" feature.
+// For more complicated schemas, it is recommended to use a dedicated JSON schema library
+// and/or pass in the schema in []byte format.
+package jsonschema
+
+import "encoding/json"
+
+type DataType string
+
+const (
+	Object  DataType = "object"
+	Number  DataType = "number"
+	Integer DataType = "integer"
+	String  DataType = "string"
+	Array   DataType = "array"
+	Null    DataType = "null"
+	Boolean DataType = "boolean"
+)
+
+// Definition is a struct for describing a JSON Schema.
+// It is fairly limited, and you may have better luck using a third-party library.
+type Definition struct {
+	// Type specifies the data type of the schema.
+	Type DataType `json:"type,omitempty"`
+	// Description is the description of the schema.
+	Description string `json:"description,omitempty"`
+	// Enum is used to restrict a value to a fixed set of values. It must be an array with at least
+	// one element, where each element is unique. You will probably only use this with strings.
+	Enum []string `json:"enum,omitempty"`
+	// Properties describes the properties of an object, if the schema type is Object.
+	Properties map[string]Definition `json:"properties"`
+	// Required specifies which properties are required, if the schema type is Object.
+	Required []string `json:"required,omitempty"`
+	// Items specifies which data type an array contains, if the schema type is Array.
+	Items *Definition `json:"items,omitempty"`
+}
+
+func (d Definition) MarshalJSON() ([]byte, error) {
+	if d.Properties == nil {
+		d.Properties = make(map[string]Definition)
+	}
+	type Alias Definition
+	return json.Marshal(struct {
+		Alias
+	}{
+		Alias: (Alias)(d),
+	})
+}


Test patch:
diff --git a/jsonschema/json_test.go b/jsonschema/json_test.go
new file mode 100644
index 00000000..716ef4ec
--- /dev/null
+++ b/jsonschema/json_test.go
@@ -0,0 +1,214 @@
+package jsonschema_test
+
+import (
+	"encoding/json"
+	"reflect"
+	"testing"
+
+	. "github.com/tmc/langchaingo/jsonschema"
+)
+
+func TestDefinition_MarshalJSON(t *testing.T) { //nolint:funlen
+	t.Parallel()
+	tests := []struct {
+		name string
+		def  Definition
+		want string
+	}{
+		{
+			name: "Test with empty Definition",
+			def:  Definition{},
+			want: `{"properties":{}}`,
+		},
+		{
+			name: "Test with Definition properties set",
+			def: Definition{
+				Type:        String,
+				Description: "A string type",
+				Properties: map[string]Definition{
+					"name": {
+						Type: String,
+					},
+				},
+			},
+			want: `{
+   "type":"string",
+   "description":"A string type",
+   "properties":{
+      "name":{
+         "type":"string",
+         "properties":{}
+      }
+   }
+}`,
+		},
+		{
+			name: "Test with nested Definition properties",
+			def: Definition{
+				Type: Object,
+				Properties: map[string]Definition{
+					"user": {
+						Type: Object,
+						Properties: map[string]Definition{
+							"name": {
+								Type: String,
+							},
+							"age": {
+								Type: Integer,
+							},
+						},
+					},
+				},
+			},
+			want: `{
+   "type":"object",
+   "properties":{
+      "user":{
+         "type":"object",
+         "properties":{
+            "name":{
+               "type":"string",
+               "properties":{}
+            },
+            "age":{
+               "type":"integer",
+               "properties":{}
+            }
+         }
+      }
+   }
+}`,
+		},
+		{
+			name: "Test with complex nested Definition",
+			def: Definition{
+				Type: Object,
+				Properties: map[string]Definition{
+					"user": {
+						Type: Object,
+						Properties: map[string]Definition{
+							"name": {
+								Type: String,
+							},
+							"age": {
+								Type: Integer,
+							},
+							"address": {
+								Type: Object,
+								Properties: map[string]Definition{
+									"city": {
+										Type: String,
+									},
+									"country": {
+										Type: String,
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			want: `{
+   "type":"object",
+   "properties":{
+      "user":{
+         "type":"object",
+         "properties":{
+            "name":{
+               "type":"string",
+               "properties":{}
+            },
+            "age":{
+               "type":"integer",
+               "properties":{}
+            },
+            "address":{
+               "type":"object",
+               "properties":{
+                  "city":{
+                     "type":"string",
+                     "properties":{}
+                  },
+                  "country":{
+                     "type":"string",
+                     "properties":{}
+                  }
+               }
+            }
+         }
+      }
+   }
+}`,
+		},
+		{
+			name: "Test with Array type Definition",
+			def: Definition{
+				Type: Array,
+				Items: &Definition{
+					Type: String,
+				},
+				Properties: map[string]Definition{
+					"name": {
+						Type: String,
+					},
+				},
+			},
+			want: `{
+   "type":"array",
+   "items":{
+      "type":"string",
+      "properties":{
+         
+      }
+   },
+   "properties":{
+      "name":{
+         "type":"string",
+         "properties":{}
+      }
+   }
+}`,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			wantBytes := []byte(tt.want)
+			var want map[string]interface{}
+			err := json.Unmarshal(wantBytes, &want)
+			if err != nil {
+				t.Errorf("Failed to Unmarshal JSON: error = %v", err)
+				return
+			}
+
+			got := structToMap(t, tt.def)
+			gotPtr := structToMap(t, &tt.def)
+
+			if !reflect.DeepEqual(got, want) {
+				t.Errorf("MarshalJSON() got = %v, want %v", got, want)
+			}
+			if !reflect.DeepEqual(gotPtr, want) {
+				t.Errorf("MarshalJSON() gotPtr = %v, want %v", gotPtr, want)
+			}
+		})
+	}
+}
+
+func structToMap(t *testing.T, v any) map[string]any {
+	t.Helper()
+	gotBytes, err := json.Marshal(v)
+	if err != nil {
+		t.Errorf("Failed to Marshal JSON: error = %v", err)
+		return nil
+	}
+
+	var got map[string]interface{}
+	err = json.Unmarshal(gotBytes, &got)
+	if err != nil {
+		t.Errorf("Failed to Unmarshal JSON: error =  %v", err)
+		return nil
+	}
+	return got
+}