From 9d977cde9b26bb7b4f9059c582d7a25308ab32ce Mon Sep 17 00:00:00 2001
From: Hazel <hazel.he@datastax.com>
Date: Wed, 10 Jul 2024 16:41:24 -0400
Subject: [PATCH] Improve error message when EGW timeout (#1255)

---
 .../embedding/gateway/EmbeddingGatewayClient.java  | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/src/main/java/io/stargate/sgv2/jsonapi/service/embedding/gateway/EmbeddingGatewayClient.java b/src/main/java/io/stargate/sgv2/jsonapi/service/embedding/gateway/EmbeddingGatewayClient.java
index 70eb37136f..fc6dd90b49 100644
--- a/src/main/java/io/stargate/sgv2/jsonapi/service/embedding/gateway/EmbeddingGatewayClient.java
+++ b/src/main/java/io/stargate/sgv2/jsonapi/service/embedding/gateway/EmbeddingGatewayClient.java
@@ -1,5 +1,7 @@
 package io.stargate.sgv2.jsonapi.service.embedding.gateway;
 
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
 import io.smallrye.mutiny.Uni;
 import io.stargate.embedding.gateway.EmbeddingGateway;
 import io.stargate.embedding.gateway.EmbeddingService;
@@ -147,9 +149,15 @@ else if (value instanceof Boolean)
             .setEmbeddingRequest(embeddingRequest)
             .setProviderContext(providerContext)
             .build();
-
-    final Uni<EmbeddingGateway.EmbeddingResponse> embeddingResponse =
-        embeddingService.embed(providerEmbedRequest);
+    Uni<EmbeddingGateway.EmbeddingResponse> embeddingResponse;
+    try {
+      embeddingResponse = embeddingService.embed(providerEmbedRequest);
+    } catch (StatusRuntimeException e) {
+      if (e.getStatus().getCode().equals(Status.Code.DEADLINE_EXCEEDED)) {
+        throw ErrorCode.EMBEDDING_PROVIDER_TIMEOUT.toApiException(e, e.getMessage());
+      }
+      throw e;
+    }
     return embeddingResponse
         .onItem()
         .transform(