-
Notifications
You must be signed in to change notification settings - Fork 22
/
document.go
53 lines (47 loc) · 1.53 KB
/
document.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
package chromem
import (
"context"
"errors"
)
// Document represents a single document.
type Document struct {
ID string
Metadata map[string]string
Embedding []float32
Content string
// ⚠️ When adding unexported fields here, consider adding a persistence struct
// version of this in [DB.Export] and [DB.Import].
}
// NewDocument creates a new document, including its embeddings.
// Metadata is optional.
// If the embeddings are not provided, they are created using the embedding function.
// You can leave the content empty if you only want to store embeddings.
// If embeddingFunc is nil, the default embedding function is used.
//
// If you want to create a document without embeddings, for example to let [Collection.AddDocuments]
// create them concurrently, you can create a document with `chromem.Document{...}`
// instead of using this constructor.
func NewDocument(ctx context.Context, id string, metadata map[string]string, embedding []float32, content string, embeddingFunc EmbeddingFunc) (Document, error) {
if id == "" {
return Document{}, errors.New("id is empty")
}
if len(embedding) == 0 && content == "" {
return Document{}, errors.New("either embedding or content must be filled")
}
if embeddingFunc == nil {
embeddingFunc = NewEmbeddingFuncDefault()
}
if len(embedding) == 0 {
var err error
embedding, err = embeddingFunc(ctx, content)
if err != nil {
return Document{}, err
}
}
return Document{
ID: id,
Metadata: metadata,
Embedding: embedding,
Content: content,
}, nil
}