Metrics with OpenTelemetry (#481)

* First approach for adding metrics to Conversations * Spotless apply * Spotless apply * LogsMetrics improved * Spotless apply
xebia-functional · Oct 18, 2023 · 3d216a3 · 3d216a3
1 parent bf8955d
commit 3d216a3
Show file tree

Hide file tree

Showing 33 changed files with 432 additions and 81 deletions.
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt
@@ -4,6 +4,7 @@ import com.xebia.functional.xef.AIError
 import com.xebia.functional.xef.llm.*
 import com.xebia.functional.xef.llm.models.functions.CFunction
 import com.xebia.functional.xef.llm.models.images.ImagesGenerationResponse
+import com.xebia.functional.xef.metrics.Metric
 import com.xebia.functional.xef.prompt.Prompt
 import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.VectorStore
@@ -17,6 +18,8 @@ interface Conversation : AutoClose, AutoCloseable {
 
   val store: VectorStore
 
+  val metric: Metric
+
   val conversationId: ConversationId?
 
   val conversation: Conversation
@@ -80,14 +83,16 @@ interface Conversation : AutoClose, AutoCloseable {
 
     operator fun invoke(
       store: VectorStore,
+      metric: Metric,
       conversationId: ConversationId? = ConversationId(UUID.generateUUID().toString())
-    ): PlatformConversation = PlatformConversation.create(store, conversationId)
+    ): PlatformConversation = PlatformConversation.create(store, metric, conversationId)
 
     @JvmSynthetic
     suspend operator fun <A> invoke(
       store: VectorStore,
+      metric: Metric,
       conversationId: ConversationId? = ConversationId(UUID.generateUUID().toString()),
       block: suspend PlatformConversation.() -> A
-    ): A = block(invoke(store, conversationId))
+    ): A = block(invoke(store, metric, conversationId))
   }
 }
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/PlatformConversation.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/PlatformConversation.kt
@@ -1,5 +1,6 @@
 package com.xebia.functional.xef.conversation
 
+import com.xebia.functional.xef.metrics.Metric
 import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.VectorStore
 
@@ -11,6 +12,7 @@ expect abstract class PlatformConversation(
   companion object {
     fun create(
       store: VectorStore,
+      metric: Metric,
       conversationId: ConversationId?,
     ): PlatformConversation
   }

diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Chat.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Chat.kt
@@ -45,22 +45,25 @@ interface Chat : LLM {
     promptMessages(prompt, scope).firstOrNull() ?: throw AIError.NoResponse()
 
   @AiDsl
-  suspend fun promptMessages(prompt: Prompt, scope: Conversation): List<String> {
-    val promptMemories = prompt.messages.toMemory(scope)
-    val adaptedPrompt = PromptCalculator.adaptPromptToConversationAndModel(prompt, scope, this@Chat)
+  suspend fun promptMessages(prompt: Prompt, scope: Conversation): List<String> =
+    scope.metric.promptSpan(scope, prompt) {
+      val promptMemories = prompt.messages.toMemory(scope)
+      val adaptedPrompt =
+        PromptCalculator.adaptPromptToConversationAndModel(prompt, scope, this@Chat)
 
-    val request =
-      ChatCompletionRequest(
-        user = adaptedPrompt.configuration.user,
-        messages = adaptedPrompt.messages,
-        n = adaptedPrompt.configuration.numberOfPredictions,
-        temperature = adaptedPrompt.configuration.temperature,
-        maxTokens = adaptedPrompt.configuration.minResponseTokens,
-      )
+      val request =
+        ChatCompletionRequest(
+          user = adaptedPrompt.configuration.user,
+          messages = adaptedPrompt.messages,
+          n = adaptedPrompt.configuration.numberOfPredictions,
+          temperature = adaptedPrompt.configuration.temperature,
+          maxTokens = adaptedPrompt.configuration.minResponseTokens,
+        )
 
-    return createChatCompletion(request)
-      .choices
-      .addChoiceToMemory(scope, promptMemories)
-      .mapNotNull { it.message?.content }
-  }
+      createChatCompletion(request)
+        .addMetrics(scope)
+        .choices
+        .addChoiceToMemory(scope, promptMemories)
+        .mapNotNull { it.message?.content }
+    }
 }
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/ChatWithFunctions.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/ChatWithFunctions.kt
@@ -64,37 +64,36 @@ interface ChatWithFunctions : LLM {
     scope: Conversation,
     function: CFunction,
     serializer: (json: String) -> A,
-  ): A {
-    val promptWithFunctions = prompt.copy(function = function)
-    val adaptedPrompt =
-      PromptCalculator.adaptPromptToConversationAndModel(
-        promptWithFunctions,
-        scope,
-        this@ChatWithFunctions
-      )
+  ): A =
+    scope.metric.promptSpan(scope, prompt) {
+      val promptWithFunctions = prompt.copy(function = function)
+      val adaptedPrompt =
+        PromptCalculator.adaptPromptToConversationAndModel(
+          promptWithFunctions,
+          scope,
+          this@ChatWithFunctions
+        )
 
-    val request =
-      FunChatCompletionRequest(
-        user = adaptedPrompt.configuration.user,
-        messages = adaptedPrompt.messages,
-        n = adaptedPrompt.configuration.numberOfPredictions,
-        temperature = adaptedPrompt.configuration.temperature,
-        maxTokens = adaptedPrompt.configuration.minResponseTokens,
-        functions = adaptedPrompt.function!!.nel(),
-        functionCall = mapOf("name" to (adaptedPrompt.function.name)),
-      )
+      val request =
+        FunChatCompletionRequest(
+          user = adaptedPrompt.configuration.user,
+          messages = adaptedPrompt.messages,
+          n = adaptedPrompt.configuration.numberOfPredictions,
+          temperature = adaptedPrompt.configuration.temperature,
+          maxTokens = adaptedPrompt.configuration.minResponseTokens,
+          functions = adaptedPrompt.function!!.nel(),
+          functionCall = mapOf("name" to (adaptedPrompt.function.name)),
+        )
 
-    return tryDeserialize(
-      serializer,
-      promptWithFunctions.configuration.maxDeserializationAttempts
-    ) {
-      val requestedMemories = prompt.messages.toMemory(scope)
-      createChatCompletionWithFunctions(request)
-        .choices
-        .addChoiceWithFunctionsToMemory(scope, requestedMemories)
-        .mapNotNull { it.message?.functionCall?.arguments }
+      tryDeserialize(serializer, promptWithFunctions.configuration.maxDeserializationAttempts) {
+        val requestedMemories = prompt.messages.toMemory(scope)
+        createChatCompletionWithFunctions(request)
+          .addMetrics(scope)
+          .choices
+          .addChoiceWithFunctionsToMemory(scope, requestedMemories)
+          .mapNotNull { it.message?.functionCall?.arguments }
+      }
     }
-  }
 
   @AiDsl
   fun <A> promptStreaming(

diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Embeddings.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Embeddings.kt
@@ -18,7 +18,9 @@ interface Embeddings : LLM {
     else
       texts
         .chunked(chunkSize ?: 400)
-        .parMap { createEmbeddings(EmbeddingRequest(name, it, requestConfig.user.id)).data }
+        .parMap {
+          createEmbeddings(EmbeddingRequest(modelType.name, it, requestConfig.user.id)).data
+        }
         .flatten()
 
   suspend fun embedQuery(text: String, requestConfig: RequestConfig): List<Embedding> =

diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/MetricManagement.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/MetricManagement.kt
@@ -0,0 +1,25 @@
+package com.xebia.functional.xef.llm
+
+import com.xebia.functional.xef.conversation.Conversation
+import com.xebia.functional.xef.llm.models.chat.ChatCompletionResponse
+import com.xebia.functional.xef.llm.models.chat.ChatCompletionResponseWithFunctions
+
+fun ChatCompletionResponseWithFunctions.addMetrics(
+  conversation: Conversation
+): ChatCompletionResponseWithFunctions {
+  conversation.metric.log(conversation, "Model: ${`object`}")
+  conversation.metric.log(
+    conversation,
+    "Tokens: ${usage.promptTokens} (prompt) + ${usage.completionTokens} (completion) = ${usage.totalTokens}"
+  )
+  return this
+}
+
+fun ChatCompletionResponse.addMetrics(conversation: Conversation): ChatCompletionResponse {
+  conversation.metric.log(conversation, "Model: ${`object`}")
+  conversation.metric.log(
+    conversation,
+    "Tokens: ${usage.promptTokens} (prompt) + ${usage.completionTokens} (completion) = ${usage.totalTokens}"
+  )
+  return this
+}
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/metrics/LogsMetric.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/metrics/LogsMetric.kt
@@ -0,0 +1,29 @@
+package com.xebia.functional.xef.metrics
+
+import com.xebia.functional.xef.conversation.Conversation
+import com.xebia.functional.xef.prompt.Prompt
+import io.ktor.util.date.*
+
+class LogsMetric : Metric {
+
+  private val identSize = 4
+
+  override suspend fun <A> promptSpan(
+    conversation: Conversation,
+    prompt: Prompt,
+    block: suspend Metric.() -> A
+  ): A {
+    val milis = getTimeMillis()
+    val name = prompt.messages.lastOrNull()?.content ?: "empty"
+    println("Prompt-Span: $name")
+    val output = block()
+    println("${writeIdent()}|-- Finished in ${getTimeMillis()-milis} ms")
+    return output
+  }
+
+  override fun log(conversation: Conversation, message: String) {
+    println("${writeIdent()}|-- $message".padStart(identSize, ' '))
+  }
+
+  private fun writeIdent(times: Int = 1) = (1..identSize * times).fold("") { a, b -> "$a " }
+}
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/metrics/Metric.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/metrics/Metric.kt
@@ -0,0 +1,14 @@
+package com.xebia.functional.xef.metrics
+
+import com.xebia.functional.xef.conversation.Conversation
+import com.xebia.functional.xef.prompt.Prompt
+
+interface Metric {
+  suspend fun <A> promptSpan(
+    conversation: Conversation,
+    prompt: Prompt,
+    block: suspend Metric.() -> A
+  ): A
+
+  fun log(conversation: Conversation, message: String)
+}
diff --git a/core/src/commonTest/kotlin/com/xebia/functional/xef/conversation/ConversationSpec.kt b/core/src/commonTest/kotlin/com/xebia/functional/xef/conversation/ConversationSpec.kt
@@ -4,6 +4,7 @@ import com.xebia.functional.tokenizer.ModelType
 import com.xebia.functional.xef.data.*
 import com.xebia.functional.xef.llm.models.chat.Message
 import com.xebia.functional.xef.llm.models.chat.Role
+import com.xebia.functional.xef.metrics.LogsMetric
 import com.xebia.functional.xef.prompt.Prompt
 import com.xebia.functional.xef.prompt.templates.assistant
 import com.xebia.functional.xef.prompt.templates.system
@@ -26,7 +27,12 @@ class ConversationSpec :
 
       val model = TestModel(modelType = ModelType.ADA)
 
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
 
       val vectorStore = scope.store
 
@@ -48,7 +54,12 @@ class ConversationSpec :
       |""" {
       val messages = generateRandomMessages(50, 40, 60)
       val conversationId = ConversationId(UUID.generateUUID().toString())
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
       val vectorStore = scope.store
 
       val modelAda = TestModel(modelType = ModelType.ADA, responses = messages)
@@ -85,7 +96,12 @@ class ConversationSpec :
       |""" {
       val messages = generateRandomMessages(50, 40, 60)
       val conversationId = ConversationId(UUID.generateUUID().toString())
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
       val vectorStore = scope.store
 
       val modelGPTTurbo16K =
@@ -122,7 +138,12 @@ class ConversationSpec :
       val message = mapOf(question to Json.encodeToString(answer))
 
       val conversationId = ConversationId(UUID.generateUUID().toString())
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
 
       val model =
         TestFunctionsModel(modelType = ModelType.GPT_3_5_TURBO_FUNCTIONS, responses = message)
@@ -146,7 +167,12 @@ class ConversationSpec :
       val message = mapOf(questionJsonString to answerJsonString)
 
       val conversationId = ConversationId(UUID.generateUUID().toString())
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
 
       val model =
         TestFunctionsModel(modelType = ModelType.GPT_3_5_TURBO_FUNCTIONS, responses = message)
@@ -170,7 +196,12 @@ class ConversationSpec :
 
       val model = TestModel(modelType = ModelType.ADA)
 
-      val scope = Conversation(LocalVectorStore(TestEmbeddings()), conversationId = conversationId)
+      val scope =
+        Conversation(
+          LocalVectorStore(TestEmbeddings()),
+          LogsMetric(),
+          conversationId = conversationId
+        )
 
       val vectorStore = scope.store
 
@@ -218,7 +249,7 @@ class ConversationSpec :
 
       val vectorStore = LocalVectorStore(TestEmbeddings())
 
-      val scope1 = Conversation(vectorStore, conversationId = conversationId)
+      val scope1 = Conversation(vectorStore, LogsMetric(), conversationId = conversationId)
 
       val firstPrompt = Prompt {
         +user("question in scope 1")
@@ -227,7 +258,7 @@ class ConversationSpec :
 
       model.promptMessages(prompt = firstPrompt, scope = scope1)
 
-      val scope2 = Conversation(vectorStore, conversationId = conversationId)
+      val scope2 = Conversation(vectorStore, LogsMetric(), conversationId = conversationId)
 
       val secondPrompt = Prompt {
         +user("question in scope 2")

diff --git a/core/src/jsMain/kotlin/com/xebia/functional/xef/conversation/JSConversation.kt b/core/src/jsMain/kotlin/com/xebia/functional/xef/conversation/JSConversation.kt
@@ -1,10 +1,12 @@
 package com.xebia.functional.xef.conversation
 
+import com.xebia.functional.xef.metrics.Metric
 import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.VectorStore
 
 class JSConversation(
   override val store: VectorStore,
+  override val metric: Metric,
   override val conversationId: ConversationId?
 ) : PlatformConversation(store, conversationId) {
 

diff --git a/core/src/jsMain/kotlin/com/xebia/functional/xef/conversation/PlatformConversation.js.kt b/core/src/jsMain/kotlin/com/xebia/functional/xef/conversation/PlatformConversation.js.kt
@@ -1,14 +1,19 @@
 package com.xebia.functional.xef.conversation
 
+import com.xebia.functional.xef.metrics.Metric
 import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.VectorStore
 
 actual abstract class PlatformConversation
 actual constructor(store: VectorStore, conversationId: ConversationId?) : Conversation, AutoClose {
   actual companion object {
-    actual fun create(store: VectorStore, conversationId: ConversationId?): PlatformConversation {
+    actual fun create(
+      store: VectorStore,
+      metric: Metric,
+      conversationId: ConversationId?
+    ): PlatformConversation {
       conversationId?.let { store.updateIndexByConversationId(conversationId) }
-      return JSConversation(store, conversationId)
+      return JSConversation(store, metric, conversationId)
     }
   }
 }
diff --git a/core/src/jvmMain/kotlin/com/xebia/functional/xef/conversation/JVMConversation.kt b/core/src/jvmMain/kotlin/com/xebia/functional/xef/conversation/JVMConversation.kt
@@ -1,11 +1,13 @@
 package com.xebia.functional.xef.conversation
 
+import com.xebia.functional.xef.metrics.Metric
 import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.VectorStore
 import java.io.Closeable
 
 open class JVMConversation(
   override val store: VectorStore,
+  override val metric: Metric,
   override val conversationId: ConversationId?,
 ) :
   PlatformConversation(store, conversationId), AutoClose by autoClose(), AutoCloseable, Closeable {