diff --git a/ai/ai21labs/v0/component_test.go b/ai/ai21labs/v0/component_test.go
index 7b1dd35dc..72da03d30 100644
--- a/ai/ai21labs/v0/component_test.go
+++ b/ai/ai21labs/v0/component_test.go
@@ -9,6 +9,7 @@ import (
 
 	qt "github.com/frankban/quicktest"
 	"github.com/instill-ai/component/base"
+	"github.com/instill-ai/component/internal/mock"
 	"github.com/instill-ai/component/internal/util/httpclient"
 	"github.com/instill-ai/x/errmsg"
 	"go.uber.org/zap"
@@ -120,8 +121,16 @@ func TestComponent_Connection(t *testing.T) {
 		})
 		c.Assert(err, qt.IsNil)
 		pbIn := new(structpb.Struct)
-		_, err = exec.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Check(err, qt.IsNil)
+
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			return nil
+		})
+
+		err = exec.Execute(ctx, ir, ow)
+		c.Assert(err, qt.IsNil)
 
 	})
 
@@ -142,8 +151,15 @@ func TestComponent_Connection(t *testing.T) {
 		})
 		c.Assert(err, qt.IsNil)
 		pbIn := new(structpb.Struct)
-		_, err = exec.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Check(err, qt.IsNotNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			return nil
+		})
+
+		err = exec.Execute(ctx, ir, ow)
+		c.Assert(err, qt.IsNotNil)
 
 		want := "AI21labs responded with a 401 status code. Incorrect API key provided."
 		c.Check(errmsg.Message(err), qt.Equals, want)
diff --git a/ai/ai21labs/v0/main.go b/ai/ai21labs/v0/main.go
index a65e15ef4..545c403ad 100644
--- a/ai/ai21labs/v0/main.go
+++ b/ai/ai21labs/v0/main.go
@@ -122,20 +122,24 @@ func (e *execution) UsesInstillCredentials() bool {
 	return e.usesInstillCredentials
 }
 
-func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(ctx context.Context, in base.InputReader, out base.OutputWriter) error {
+	inputs, err := in.Read(ctx)
+	if err != nil {
+		return err
+	}
 	outputs := make([]*structpb.Struct, len(inputs))
 
 	// The execution takes a array of inputs and returns an array of outputs. The execution is done sequentially.
 	for i, input := range inputs {
 		output, err := e.execute(input)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		outputs[i] = output
 	}
 
-	return outputs, nil
+	return out.Write(ctx, outputs)
 }
 
 // This function is not used in the codebase. It is only used in the tests.
diff --git a/ai/ai21labs/v0/tasks_test.go b/ai/ai21labs/v0/tasks_test.go
index eabd76017..fa7eb6ded 100644
--- a/ai/ai21labs/v0/tasks_test.go
+++ b/ai/ai21labs/v0/tasks_test.go
@@ -9,6 +9,7 @@ import (
 	qt "github.com/frankban/quicktest"
 	"github.com/instill-ai/component/ai"
 	"github.com/instill-ai/component/base"
+	"github.com/instill-ai/component/internal/mock"
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/types/known/structpb"
 )
@@ -184,12 +185,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task embedding", func(c *qt.C) {
@@ -219,12 +226,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task contextual answers without answers", func(c *qt.C) {
@@ -253,12 +266,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 
 	})
 
@@ -288,12 +307,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task text summarization", func(c *qt.C) {
@@ -320,12 +345,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task text summarization by segment", func(c *qt.C) {
@@ -358,12 +389,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task text paraphrasing", func(c *qt.C) {
@@ -389,12 +426,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task grammar check", func(c *qt.C) {
@@ -429,12 +472,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task text improvement", func(c *qt.C) {
@@ -468,12 +517,18 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 
 	c.Run("ok - task text segmentation", func(c *qt.C) {
@@ -505,11 +560,17 @@ func TestTasks(t *testing.T) {
 		pbIn, err := base.ConvertToStructpb(tc.input)
 		c.Assert(err, qt.IsNil)
 
-		got, err := e.Execute(ctx, []*structpb.Struct{pbIn})
-		c.Assert(err, qt.IsNil)
+		ir := mock.NewInputReaderMock(c)
+		ow := mock.NewOutputWriterMock(c)
+		ir.ReadMock.Return([]*structpb.Struct{pbIn}, nil)
+		ow.WriteMock.Optional().Set(func(ctx context.Context, outputs []*structpb.Struct) (err error) {
+			wantJSON, err := json.Marshal(tc.wantResp)
+			c.Assert(err, qt.IsNil)
+			c.Check(wantJSON, qt.JSONEquals, outputs[0].AsMap())
+			return nil
+		})
 
-		wantJSON, err := json.Marshal(tc.wantResp)
+		err = e.Execute(ctx, ir, ow)
 		c.Assert(err, qt.IsNil)
-		c.Check(wantJSON, qt.JSONEquals, got[0].AsMap())
 	})
 }