From fedc9f73e6272d6ce4420e34744928ae319a72ed Mon Sep 17 00:00:00 2001 From: voocel Date: Wed, 4 Jun 2025 23:07:00 +0800 Subject: [PATCH] feat(schema): add utility methods for Document --- schema/document.go | 66 +++++++++++++++++++++++ schema/document_test.go | 112 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 178 insertions(+) diff --git a/schema/document.go b/schema/document.go index b0b1c2b1..3027b585 100644 --- a/schema/document.go +++ b/schema/document.go @@ -201,3 +201,69 @@ func (d *Document) SparseVector() map[int]float64 { return nil } + +// Clone creates a deep copy of the document. +// This is useful when you need to create a copy of the document without affecting the original. +func (d *Document) Clone() *Document { + if d == nil { + return nil + } + + clone := &Document{ + ID: d.ID, + Content: d.Content, + } + + if d.MetaData != nil { + clone.MetaData = make(map[string]any, len(d.MetaData)) + for k, v := range d.MetaData { + // For simplicity, we do a shallow copy of the values. + // For deep copy of complex types, users can implement their own logic. + clone.MetaData[k] = v + } + } + + return clone +} + +// IsEmpty checks if the document is empty (no ID, no content, and no metadata). +func (d *Document) IsEmpty() bool { + return d.ID == "" && d.Content == "" && len(d.MetaData) == 0 +} + +// HasMetaData checks if the document contains the specified metadata key. +func (d *Document) HasMetaData(key string) bool { + if d.MetaData == nil { + return false + } + _, exists := d.MetaData[key] + return exists +} + +// ClearMetaData removes all metadata from the document. +// Returns the document itself for method chaining. +func (d *Document) ClearMetaData() *Document { + d.MetaData = nil + return d +} + +// WithMetaData sets custom metadata for the document. +// This is a generic method for setting any metadata key-value pair. +// Returns the document itself for method chaining. +func (d *Document) WithMetaData(key string, value any) *Document { + if d.MetaData == nil { + d.MetaData = make(map[string]any) + } + d.MetaData[key] = value + return d +} + +// GetMetaData retrieves custom metadata from the document. +// Returns the value and a boolean indicating whether the key exists. +func (d *Document) GetMetaData(key string) (any, bool) { + if d.MetaData == nil { + return nil, false + } + value, exists := d.MetaData[key] + return value, exists +} diff --git a/schema/document_test.go b/schema/document_test.go index db2ba374..1ca75cad 100644 --- a/schema/document_test.go +++ b/schema/document_test.go @@ -51,3 +51,115 @@ func TestDocument(t *testing.T) { convey.So(d.DenseVector(), convey.ShouldEqual, vector) }) } + +func TestDocumentUtilityMethods(t *testing.T) { + convey.Convey("test document utility methods", t, func() { + convey.Convey("test Clone method", func() { + original := &Document{ + ID: "test-id", + Content: "test content", + MetaData: map[string]any{"key1": "value1", "key2": 42}, + } + + clone := original.Clone() + + convey.So(clone, convey.ShouldNotBeNil) + convey.So(clone.ID, convey.ShouldEqual, original.ID) + convey.So(clone.Content, convey.ShouldEqual, original.Content) + convey.So(clone.MetaData["key1"], convey.ShouldEqual, "value1") + convey.So(clone.MetaData["key2"], convey.ShouldEqual, 42) + + // Test independence: modifying clone doesn't affect original + clone.MetaData["key3"] = "new value" + convey.So(original.MetaData["key3"], convey.ShouldBeNil) + + // Test edge cases + var nilDoc *Document + convey.So(nilDoc.Clone(), convey.ShouldBeNil) + + docWithNilMeta := &Document{ID: "test", Content: "test"} + convey.So(docWithNilMeta.Clone().MetaData, convey.ShouldBeNil) + }) + + convey.Convey("test IsEmpty method", func() { + convey.So((&Document{}).IsEmpty(), convey.ShouldBeTrue) + convey.So((&Document{MetaData: map[string]any{}}).IsEmpty(), convey.ShouldBeTrue) + + convey.So((&Document{ID: "test"}).IsEmpty(), convey.ShouldBeFalse) + convey.So((&Document{Content: "test"}).IsEmpty(), convey.ShouldBeFalse) + convey.So((&Document{MetaData: map[string]any{"key": "value"}}).IsEmpty(), convey.ShouldBeFalse) + }) + + convey.Convey("test HasMetaData method", func() { + doc := &Document{ + MetaData: map[string]any{"existing_key": "value", "nil_value": nil}, + } + + convey.So(doc.HasMetaData("existing_key"), convey.ShouldBeTrue) + convey.So(doc.HasMetaData("nil_value"), convey.ShouldBeTrue) // nil value still counts as existing + convey.So(doc.HasMetaData("non_existing"), convey.ShouldBeFalse) + + // Test with nil metadata + convey.So((&Document{}).HasMetaData("any_key"), convey.ShouldBeFalse) + }) + + convey.Convey("test ClearMetaData method", func() { + doc := &Document{ + ID: "test-id", + Content: "test content", + MetaData: map[string]any{"key1": "value1", "key2": "value2"}, + } + + result := doc.ClearMetaData() + + convey.So(result, convey.ShouldEqual, doc) // Should return same instance for chaining + convey.So(doc.MetaData, convey.ShouldBeNil) + convey.So(doc.ID, convey.ShouldEqual, "test-id") // Other fields unchanged + }) + + convey.Convey("test WithMetaData and GetMetaData methods", func() { + doc := &Document{ID: "test-id"} + + // Test setting and getting metadata + result := doc.WithMetaData("custom_key", "custom_value") + convey.So(result, convey.ShouldEqual, doc) // Should return same instance + convey.So(doc.MetaData, convey.ShouldNotBeNil) + + value, exists := doc.GetMetaData("custom_key") + convey.So(exists, convey.ShouldBeTrue) + convey.So(value, convey.ShouldEqual, "custom_value") + + // Test different value types + doc.WithMetaData("int_key", 42).WithMetaData("bool_key", true) + + intValue, intExists := doc.GetMetaData("int_key") + convey.So(intExists, convey.ShouldBeTrue) + convey.So(intValue, convey.ShouldEqual, 42) + + // Test non-existing key + _, nonExists := doc.GetMetaData("non_existing") + convey.So(nonExists, convey.ShouldBeFalse) + + // Test GetMetaData with nil MetaData + nilValue, nilExists := (&Document{}).GetMetaData("any_key") + convey.So(nilExists, convey.ShouldBeFalse) + convey.So(nilValue, convey.ShouldBeNil) + }) + + convey.Convey("test method chaining", func() { + doc := &Document{} + + result := doc.WithMetaData("key1", "value1"). + WithScore(0.95). + WithExtraInfo("test info") + + convey.So(result, convey.ShouldEqual, doc) + convey.So(doc.Score(), convey.ShouldEqual, 0.95) + convey.So(doc.ExtraInfo(), convey.ShouldEqual, "test info") + + value1, exists1 := doc.GetMetaData("key1") + convey.So(exists1, convey.ShouldBeTrue) + convey.So(value1, convey.ShouldEqual, "value1") + }) + }) +}