From b134df30b8a6952281c2f4b1a4be3d522cc17419 Mon Sep 17 00:00:00 2001 From: Mattt Zmuda Date: Wed, 2 Oct 2024 09:18:48 -0700 Subject: [PATCH] Fix behavior of run when wait is set --- run.go | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/run.go b/run.go index 975fc12..7c13865 100644 --- a/run.go +++ b/run.go @@ -82,10 +82,14 @@ func (r *Client) RunWithOptions(ctx context.Context, identifier string, input Pr return nil, err } - // Wait for the prediction to complete - err = r.Wait(ctx, prediction) - if err != nil { - return nil, err + // Check if the prediction is done based on blocking preference and status + isDone := options.blockUntilDone && prediction.Status != Starting + if !isDone { + // Wait for the prediction to complete + err = r.Wait(ctx, prediction) + if err != nil { + return nil, err + } } // Check for model error in the prediction