From 88dad4f944a86d669b73a152cf972482d3d76065 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Gill=C3=A9?= <philipp.gille@gmail.com>
Date: Sun, 1 Sep 2024 18:02:12 +0200
Subject: [PATCH] Implement Collection.GetByID()

---
 collection.go | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/collection.go b/collection.go
index ec641a1..8a0bdcb 100644
--- a/collection.go
+++ b/collection.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"maps"
 	"path/filepath"
 	"slices"
 	"sync"
@@ -291,6 +292,31 @@ func (c *Collection) AddDocument(ctx context.Context, doc Document) error {
 	return nil
 }
 
+// GetByID returns a document by its ID.
+// The returned document is a copy of the original document, so it can be safely
+// modified without affecting the collection.
+func (c *Collection) GetByID(ctx context.Context, id string) (Document, error) {
+	if id == "" {
+		return Document{}, errors.New("document ID is empty")
+	}
+
+	c.documentsLock.RLock()
+	defer c.documentsLock.RUnlock()
+
+	doc, ok := c.documents[id]
+	if ok {
+		// Clone the document
+		res := *doc
+		// Above copies the simple fields, but we need to copy the slices and maps
+		res.Metadata = maps.Clone(doc.Metadata)
+		res.Embedding = slices.Clone(doc.Embedding)
+
+		return res, nil
+	}
+
+	return Document{}, fmt.Errorf("document with ID '%v' not found", id)
+}
+
 // Delete removes document(s) from the collection.
 //
 //   - where: Conditional filtering on metadata. Optional.