Merge pull request #27 from philippgille/add-collection-count

Add Collection.Count() method

collection.go

@@ -93,6 +93,13 @@ func (c *Collection) AddConcurrently(ctx context.Context, ids []string, embeddin
 	return c.add(ctx, ids, documents, embeddings, metadatas, concurrency)
 }
 
+// Count returns the number of documents in the collection.
+func (c *Collection) Count() int {
+	c.documentsLock.RLock()
+	defer c.documentsLock.RUnlock()
+	return len(c.documents)
+}
+
 // Performs a nearest neighbors query on a collection specified by UUID.
 //
 //   - queryText: The text to search for.

collection_test.go

@@ -34,3 +34,34 @@ func TestCollection_Add(t *testing.T) {
 
 	// TODO: Check expectations when documents become accessible
 }
+
+func TestCollection_Count(t *testing.T) {
+	// Create collection
+	db := chromem.NewDB()
+	name := "test"
+	metadata := map[string]string{"foo": "bar"}
+	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
+		return []float32{-0.1, 0.1, 0.2}, nil
+	}
+	c, err := db.CreateCollection(name, metadata, embeddingFunc)
+	if err != nil {
+		t.Error("expected no error, got", err)
+	}
+	if c == nil {
+		t.Error("expected collection, got nil")
+	}
+
+	// Add documents
+	ids := []string{"1", "2"}
+	metadatas := []map[string]string{{"foo": "bar"}, {"a": "b"}}
+	documents := []string{"hello world", "hallo welt"}
+	err = c.Add(context.Background(), ids, nil, metadatas, documents)
+	if err != nil {
+		t.Error("expected nil, got", err)
+	}
+
+	// Check count
+	if c.Count() != 2 {
+		t.Error("expected 2, got", c.Count())
+	}
+}