feat: add output collection metadata and prevent attaching to output collections

- Add chroma:attached_function_id metadata when creating output collections
- Add check in AttachFunction to prevent attaching functions to output collections
- Add ErrCannotAttachToOutputCollection error

Author: tanujnay112

go/pkg/common/constants.go

@@ -3,4 +3,8 @@ package common
 const (
 	DefaultTenant   = "default_tenant"
 	DefaultDatabase = "default_database"
+
+	// SourceAttachedFunctionIDKey is the metadata key used to mark output collections
+	// and link them to their attached function.
+	SourceAttachedFunctionIDKey = "chroma:source_attached_function_id"
 )

go/pkg/common/errors.go

@@ -49,11 +49,12 @@ var (
 	ErrUnknownSegmentMetadataType = errors.New("segment metadata value type not supported")
 
 	// AttachedFunction errors
-	ErrAttachedFunctionAlreadyExists = errors.New("the attached function that was being created already exists for this collection")
-	ErrAttachedFunctionNotFound      = errors.New("the requested attached function was not found")
-	ErrAttachedFunctionNotReady      = errors.New("the requested attached function exists but is still initializing")
-	ErrInvalidAttachedFunctionName   = errors.New("attached function name cannot start with reserved prefix '_deleted_'")
-	ErrHeapServiceNotEnabled         = errors.New("heap service is not enabled")
+	ErrAttachedFunctionAlreadyExists  = errors.New("the attached function that was being created already exists for this collection")
+	ErrAttachedFunctionNotFound       = errors.New("the requested attached function was not found")
+	ErrAttachedFunctionNotReady       = errors.New("the requested attached function exists but is still initializing")
+	ErrInvalidAttachedFunctionName    = errors.New("attached function name cannot start with reserved prefix '_deleted_'")
+	ErrHeapServiceNotEnabled          = errors.New("heap service is not enabled")
+	ErrCannotAttachToOutputCollection = errors.New("cannot attach function to an output collection")
 
 	// Function errors
 	ErrFunctionNotFound = errors.New("function not found")

go/pkg/sysdb/coordinator/task.go

@@ -158,6 +158,16 @@ func (s *Coordinator) AttachFunction(ctx context.Context, req *coordinatorpb.Att
 			return common.ErrCollectionNotFound
 		}
 
+		// Check if input collection is an output collection (prevent chaining)
+		inputCollection := collections[0]
+		for _, meta := range inputCollection.CollectionMetadata {
+			if meta.Key != nil && *meta.Key == common.SourceAttachedFunctionIDKey {
+				log.Error("AttachFunction: cannot attach function to an output collection",
+					zap.String("input_collection_id", req.InputCollectionId))
+				return common.ErrCannotAttachToOutputCollection
+			}
+		}
+
 		// Serialize params
 		var paramsJSON string
 		if req.Params != nil {
@@ -574,14 +584,19 @@ func (s *Coordinator) FinishCreateAttachedFunction(ctx context.Context, req *coo
 
 		// 5. Create the output collection with segments
 		dimension := int32(1) // Default dimension for attached function output collections
+
+		// Add metadata to mark this as an output collection
+		outputCollectionMetadata := model.NewCollectionMetadata[model.CollectionMetadataValueType]()
+		outputCollectionMetadata.Add(common.SourceAttachedFunctionIDKey, &model.CollectionMetadataValueStringType{Value: attachedFunctionID.String()})
+
 		collection := &model.CreateCollection{
 			ID:                   collectionID,
 			Name:                 attachedFunction.OutputCollectionName,
 			ConfigurationJsonStr: "{}", // Empty JSON object for default config
 			TenantID:             attachedFunction.TenantID,
 			DatabaseName:         database.Name,
 			Dimension:            &dimension,
-			Metadata:             nil,
+			Metadata:             outputCollectionMetadata,
 		}
 
 		// Create segments for the collection (distributed setup)