Chat completion intercept adapter

Author: peterbanda

openai-core/src/main/scala/io/cequence/openaiscala/domain/ChatCompletionInterceptData.scala

@@ -0,0 +1,12 @@
+package io.cequence.openaiscala.domain
+
+import io.cequence.openaiscala.domain.response.ChatCompletionResponse
+import io.cequence.openaiscala.domain.settings.CreateChatCompletionSettings
+
+case class ChatCompletionInterceptData(
+  messages: Seq[BaseMessage],
+  setting: CreateChatCompletionSettings,
+  response: ChatCompletionResponse,
+  timeRequestReceived: java.util.Date,
+  timeResponseReceived: java.util.Date
+)

openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/ChatCompletionInterceptAdapter.scala

@@ -0,0 +1,58 @@
+package io.cequence.openaiscala.service.adapter
+
+import io.cequence.openaiscala.domain.{BaseMessage, ChatCompletionInterceptData}
+import io.cequence.openaiscala.domain.response.ChatCompletionResponse
+import io.cequence.openaiscala.domain.settings.CreateChatCompletionSettings
+import io.cequence.openaiscala.service.OpenAIChatCompletionService
+import io.cequence.wsclient.service.CloseableService
+import io.cequence.wsclient.service.adapter.ServiceWrapper
+
+import scala.concurrent.{ExecutionContext, Future}
+
+private class ChatCompletionInterceptAdapter[S <: OpenAIChatCompletionService](
+  intercept: ChatCompletionInterceptData => Future[Unit]
+)(
+  underlying: S
+)(
+  implicit ec: ExecutionContext
+) extends ServiceWrapper[S]
+    with CloseableService
+    with OpenAIChatCompletionService {
+
+  // we just delegate all the calls to the underlying service
+  override def wrap[T](
+    fun: S => Future[T]
+  ): Future[T] = fun(underlying)
+
+  // but for the chat completion we adapt the messages and settings
+  override def createChatCompletion(
+    messages: Seq[BaseMessage],
+    settings: CreateChatCompletionSettings
+  ): Future[ChatCompletionResponse] = {
+    val timeRequestReceived = new java.util.Date()
+
+    for {
+      response <- underlying.createChatCompletion(
+        messages,
+        settings
+      )
+
+      _ <- {
+        val timeResponseReceived = new java.util.Date()
+
+        intercept(
+          ChatCompletionInterceptData(
+            messages,
+            settings,
+            response,
+            timeRequestReceived,
+            timeResponseReceived
+          )
+        )
+      }
+    } yield response
+  }
+
+  override def close(): Unit =
+    underlying.close()
+}

openai-core/src/main/scala/io/cequence/openaiscala/service/adapter/OpenAIServiceAdapters.scala

@@ -1,13 +1,13 @@
 package io.cequence.openaiscala.service.adapter
 
-import io.cequence.openaiscala.domain.BaseMessage
+import io.cequence.openaiscala.domain.{BaseMessage, ChatCompletionInterceptData}
 import io.cequence.openaiscala.domain.settings.CreateChatCompletionSettings
 import io.cequence.openaiscala.service._
 import io.cequence.openaiscala.service.adapter.ServiceWrapperTypes._
 import io.cequence.wsclient.service.CloseableService
 import io.cequence.wsclient.service.adapter.ServiceWrapperTypes.CloseableServiceWrapper
 
-import scala.concurrent.ExecutionContext
+import scala.concurrent.{ExecutionContext, Future}
 
 object OpenAIServiceAdapters {
 
@@ -41,6 +41,17 @@ trait OpenAIServiceAdapters[S <: CloseableService] extends ServiceAdapters[S] {
       new ChatCompletionInputAdapter(adaptMessages, adaptSettings)(service)
     )
 
+  def chatCompletionIntercept(
+    intercept: ChatCompletionInterceptData => Future[Unit]
+  )(
+    service: S with OpenAIChatCompletionService
+  )(
+    implicit ec: ExecutionContext
+  ): S =
+    wrapAndDelegateChatCompletion(
+      new ChatCompletionInterceptAdapter(intercept)(service)
+    )
+
   def chatCompletionRouter(
     serviceModels: Map[OpenAIChatCompletionService, Seq[String]],
     service: S with OpenAIChatCompletionService