Add DB.GetCollection()

db.go

@@ -62,3 +62,41 @@ func (c *DB) ListCollections() map[string]*Collection {
 
 	return res
 }
+
+// GetCollection returns the collection with the given name.
+// The collection is a copy of the original, except for its EmbeddingFunc.
+// It's only copied one level deep though. So while you *can* manipulate the collection's
+// map of documents, you must not manipulate the documents themselves.
+// Regarding the EmbeddingFunc it's the original. So if it closes over some state,
+// this state is shared. But usually an EmbeddingFunc just closes over an API key
+// or HTTP client, which are safe to share.
+func (c *DB) GetCollection(name string) *Collection {
+	c.collectionsLock.RLock()
+	defer c.collectionsLock.RUnlock()
+
+	orig, ok := c.collections[name]
+	if !ok {
+		return nil
+	}
+
+	newMetadata := make(map[string]string, len(orig.Metadata))
+	for k, v := range orig.Metadata {
+		newMetadata[k] = v
+	}
+
+	orig.documentsLock.RLock()
+	defer orig.documentsLock.RUnlock()
+	newDocuments := make(map[string]*document, len(orig.documents))
+	for k, v := range orig.documents {
+		newDocuments[k] = v
+	}
+
+	return &Collection{
+		Name:     orig.Name,
+		Metadata: newMetadata,
+
+		documents: make(map[string]*document, len(orig.documents)),
+
+		embed: orig.embed,
+	}
+}

db_test.go

@@ -47,3 +47,34 @@ func TestDB_ListCollections(t *testing.T) {
 		t.Error("expected 1 collection, got", len(db.ListCollections()))
 	}
 }
+
+func TestDB_GetCollection(t *testing.T) {
+	// Values in the collection
+	name := "test"
+	metadata := map[string]string{"foo": "bar"}
+	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
+		return []float32{-0.1, 0.1, 0.2}, nil
+	}
+
+	// Create initial collection
+	db := chromem.NewDB()
+	// We ignore the return value. CreateCollection is tested elsewhere.
+	_ = db.CreateCollection(name, metadata, embeddingFunc)
+
+	// Get collection
+	c := db.GetCollection(name)
+
+	// Check expectations
+	if c.Name != name {
+		t.Error("expected name", name, "got", c.Name)
+	}
+	if len(c.Metadata) != 1 {
+		t.Error("expected 1 metadata, got", len(c.Metadata))
+	}
+	if c.Metadata["foo"] != "bar" {
+		t.Error("expected metadata", metadata, "got", c.Metadata)
+	}
+	// TODO: Check documents content as soon as we have access to them
+	// TODO: Same for the EmbeddingFunc
+	// TODO: Check documents map being a copy as soon as we have access to it
+}