From 8b3a581362397257b3af078ee44a78c07b0c1a50 Mon Sep 17 00:00:00 2001
From: Peter Banda <peter.banda@protonmail.com>
Date: Tue, 12 Nov 2024 13:28:25 +0100
Subject: [PATCH] Retry adapter - isRetryable exposed

---
 build.sbt                                                | 2 +-
 .../service/adapter/OpenAIServiceAdapters.scala          | 9 +++++++--
 .../service/adapter/RetryServiceAdapter.scala            | 8 +++++---
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/build.sbt b/build.sbt
index 56f1a26e..419412e0 100755
--- a/build.sbt
+++ b/build.sbt
@@ -7,7 +7,7 @@ val scala3 = "3.2.2"
 
 ThisBuild / organization := "io.cequence"
 ThisBuild / scalaVersion := scala212
-ThisBuild / version := "1.1.1.RC.9"
+ThisBuild / version := "1.1.1.RC.11"
 ThisBuild / isSnapshot := false
 
 lazy val commonSettings = Seq(
diff --git a/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/OpenAIServiceAdapters.scala b/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/OpenAIServiceAdapters.scala
index 2f468408..6c8c876a 100644
--- a/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/OpenAIServiceAdapters.scala
+++ b/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/OpenAIServiceAdapters.scala
@@ -3,6 +3,7 @@ package io.cequence.openaiscala.service.adapter
 import akka.actor.Scheduler
 import akka.stream.Materializer
 import io.cequence.openaiscala.RetryHelpers.RetrySettings
+import io.cequence.openaiscala.Retryable
 import io.cequence.openaiscala.domain.BaseMessage
 import io.cequence.openaiscala.domain.settings.CreateChatCompletionSettings
 import io.cequence.openaiscala.service._
@@ -44,13 +45,17 @@ trait OpenAIServiceAdapters[S <: CloseableService] {
 
   def retry(
     underlying: S,
-    log: Option[String => Unit] = None
+    log: Option[String => Unit] = None,
+    isRetryable: Throwable => Boolean = {
+      case Retryable(_) => true
+      case _            => false
+    }
   )(
     implicit ec: ExecutionContext,
     retrySettings: RetrySettings,
     scheduler: Scheduler
   ): S =
-    wrapAndDelegate(new RetryServiceAdapter(underlying, log))
+    wrapAndDelegate(new RetryServiceAdapter(underlying, log, isRetryable))
 
   def log(
     underlying: S,
diff --git a/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/RetryServiceAdapter.scala b/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/RetryServiceAdapter.scala
index a51b636e..8fbec409 100644
--- a/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/RetryServiceAdapter.scala
+++ b/openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/RetryServiceAdapter.scala
@@ -1,7 +1,7 @@
 package io.cequence.openaiscala.service.adapter
 
 import akka.actor.Scheduler
-import io.cequence.openaiscala.RetryHelpers
+import io.cequence.openaiscala.{RetryHelpers, Retryable}
 import io.cequence.openaiscala.RetryHelpers.RetrySettings
 import io.cequence.wsclient.service.CloseableService
 
@@ -9,7 +9,8 @@ import scala.concurrent.{ExecutionContext, Future}
 
 private class RetryServiceAdapter[+S <: CloseableService](
   underlying: S,
-  log: Option[String => Unit] = None
+  log: Option[String => Unit] = None,
+  isRetryable: Throwable => Boolean
 )(
   implicit ec: ExecutionContext,
   retrySettings: RetrySettings,
@@ -24,7 +25,8 @@ private class RetryServiceAdapter[+S <: CloseableService](
   ): Future[T] =
     fun(underlying).retryOnFailure(
       Some(s"${getFunctionName().capitalize} call failed"),
-      log
+      log,
+      isRetryable
     )
 
   override def close(): Unit =