New PromptBuilder returning list of messages (#326)

* New PromptBuilder returning list of messages

* Comments addressed

---------

Co-authored-by: Raúl Raja Martínez <[email protected]>

core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Chat.kt

@@ -9,6 +9,9 @@ import com.xebia.functional.xef.auto.PromptConfiguration
 import com.xebia.functional.xef.llm.models.chat.*
 import com.xebia.functional.xef.llm.models.functions.CFunction
 import com.xebia.functional.xef.prompt.Prompt
+import com.xebia.functional.xef.prompt.buildPrompt
+import com.xebia.functional.xef.prompt.templates.assistant
+import com.xebia.functional.xef.prompt.templates.user
 import com.xebia.functional.xef.vectorstores.Memory
 import io.ktor.util.date.*
 import kotlinx.coroutines.flow.Flow
@@ -83,6 +86,15 @@ interface Chat : LLM {
     promptConfiguration: PromptConfiguration = PromptConfiguration.DEFAULTS
   ): List<String> = promptMessages(Prompt(question), scope, functions, promptConfiguration)
 
+  @AiDsl
+  suspend fun promptMessage(
+    prompt: Prompt,
+    scope: Conversation,
+    promptConfiguration: PromptConfiguration = PromptConfiguration.DEFAULTS
+  ): String =
+    promptMessages(prompt, scope, emptyList(), promptConfiguration).firstOrNull()
+      ?: throw AIError.NoResponse()
+
   @AiDsl
   suspend fun promptMessages(
     prompt: Prompt,
@@ -160,7 +172,7 @@ interface Chat : LLM {
 
   suspend fun String.toMessages(): List<Message> = Prompt(this).toMessages()
 
-  suspend fun Prompt.toMessages(): List<Message> = listOf(Message.userMessage { message })
+  suspend fun Prompt.toMessages(): List<Message> = buildPrompt { +user(message) }
 
   private suspend fun addMemoriesAfterStream(
     request: ChatCompletionRequest,
@@ -332,7 +344,7 @@ interface Chat : LLM {
 
         val ctxTruncated: String = modelType.encoding.truncateText(ctx, maxContextTokens)
 
-        listOf(Message.assistantMessage { ctxTruncated })
+        buildPrompt { +assistant(ctxTruncated) }
       } else {
         emptyList()
       }

core/src/commonMain/kotlin/com/xebia/functional/xef/llm/models/chat/Message.kt

@@ -1,14 +1,10 @@
 package com.xebia.functional.xef.llm.models.chat
 
+import kotlin.jvm.JvmStatic
+
 data class Message(val role: Role, val content: String, val name: String) {
   companion object {
-    suspend fun systemMessage(message: suspend () -> String) =
-      Message(role = Role.SYSTEM, content = message(), name = Role.SYSTEM.name)
-
-    suspend fun userMessage(message: suspend () -> String) =
-      Message(role = Role.USER, content = message(), name = Role.USER.name)
-
-    suspend fun assistantMessage(message: suspend () -> String) =
-      Message(role = Role.ASSISTANT, content = message(), name = Role.ASSISTANT.name)
+    @JvmStatic
+    fun apply(role: Role, content: String, name: String): Message = Message(role, content, name)
   }
 }

core/src/commonMain/kotlin/com/xebia/functional/xef/prompt/PromptBuilder.kt

@@ -0,0 +1,25 @@
+package com.xebia.functional.xef.prompt
+
+import com.xebia.functional.xef.llm.models.chat.Message
+import com.xebia.functional.xef.llm.models.chat.Role
+
+open class PromptBuilder {
+  private val items = mutableListOf<Message>()
+
+  operator fun Message.unaryPlus() {
+    items.add(this)
+  }
+
+  operator fun List<Message>.unaryPlus() {
+    items.addAll(this)
+  }
+
+  protected open fun preprocess(elements: List<Message>): List<Message> = elements
+
+  fun build(): List<Message> = preprocess(items)
+}
+
+fun String.message(role: Role): Message = Message(role, this, role.name)
+
+fun buildPrompt(block: PromptBuilder.() -> Unit): List<Message> =
+  PromptBuilder().apply { block() }.build()

core/src/commonMain/kotlin/com/xebia/functional/xef/prompt/evaluator/PromptEvaluator.kt

@@ -3,7 +3,9 @@ package com.xebia.functional.xef.prompt.evaluator
 import com.xebia.functional.xef.auto.Conversation
 import com.xebia.functional.xef.auto.PromptConfiguration
 import com.xebia.functional.xef.llm.Chat
-import com.xebia.functional.xef.llm.models.chat.Message
+import com.xebia.functional.xef.prompt.buildPrompt
+import com.xebia.functional.xef.prompt.templates.system
+import com.xebia.functional.xef.prompt.templates.user
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.builtins.MapSerializer
 import kotlinx.serialization.builtins.serializer
@@ -162,12 +164,12 @@ ${scoreConfig.joinToString("\n") { printReturn(it) }}
     val result: List<String> =
       model.promptMessages(
         messages =
-          listOf(
-            Message.systemMessage { message },
-            Message.userMessage { "Set Prompt = $prompt" },
-            Message.userMessage { "Set Response = $response" },
-            Message.userMessage { "Evaluate(Prompt, Response)" }
-          ),
+          buildPrompt {
+            +system(message)
+            +user("Set Prompt = $prompt")
+            +user("Set Response = $response")
+            +user("Evaluate(Prompt, Response)")
+          },
         scope = conversation,
         promptConfiguration = promptConfiguration,
       )

core/src/commonMain/kotlin/com/xebia/functional/xef/prompt/expressions/Expression.kt

@@ -4,6 +4,9 @@ import com.xebia.functional.xef.auto.Conversation
 import com.xebia.functional.xef.auto.PromptConfiguration
 import com.xebia.functional.xef.llm.ChatWithFunctions
 import com.xebia.functional.xef.llm.models.chat.Message
+import com.xebia.functional.xef.prompt.buildPrompt
+import com.xebia.functional.xef.prompt.templates.assistant
+import com.xebia.functional.xef.prompt.templates.system
 import io.github.oshai.kotlinlogging.KLogger
 import io.github.oshai.kotlinlogging.KotlinLogging
 
@@ -19,16 +22,8 @@ class Expression(
 
   private val generationKeys: MutableList<String> = mutableListOf()
 
-  suspend fun system(message: suspend () -> String) {
-    messages.add(Message.systemMessage(message))
-  }
-
-  suspend fun user(message: suspend () -> String) {
-    messages.add(Message.userMessage(message))
-  }
-
-  suspend fun assistant(message: suspend () -> String) {
-    messages.add(Message.assistantMessage(message))
+  fun addMessages(newMessages: List<Message>) {
+    messages.addAll(newMessages)
   }
 
   fun prompt(key: String): String {
@@ -40,14 +35,12 @@ class Expression(
     promptConfiguration: PromptConfiguration = PromptConfiguration.DEFAULTS
   ): ExpressionResult {
     block()
-    val prelude =
-      listOf(
-        Message.systemMessage { "You are an expert in replacing variables in templates" },
-      )
-    val instructionMessages =
-      listOf(
-        Message.assistantMessage { "I will replace all placeholders in the message" },
-      )
+    val prelude = buildPrompt { +system("You are an expert in replacing variables in templates") }
+
+    val instructionMessages = buildPrompt {
+      +assistant("I will replace all placeholders in the message")
+    }
+
     val values: ReplacedValues =
       model.prompt(
         messages = prelude + messages + instructionMessages,

core/src/commonMain/kotlin/com/xebia/functional/xef/prompt/templates/templates.kt

@@ -1,49 +1,56 @@
 package com.xebia.functional.xef.prompt.templates
 
-import com.xebia.functional.xef.prompt.Prompt
+import com.xebia.functional.xef.llm.models.chat.Message
+import com.xebia.functional.xef.llm.models.chat.Role
 import com.xebia.functional.xef.prompt.PromptBuilder
-import com.xebia.functional.xef.prompt.prompt
+import com.xebia.functional.xef.prompt.message
 
-fun youAre(role: String, talkingTo: String): Prompt =
-  "You are a $role talking with a $talkingTo".prompt()
+fun system(context: String): Message = context.message(Role.SYSTEM)
 
-class StepsBuilder : PromptBuilder() {
-  override fun preprocess(elements: List<Prompt>): List<Prompt> =
-    elements.mapIndexed { ix, elt -> Prompt("${ix + 1} - ${elt.message}") }
+fun assistant(context: String): Message = context.message(Role.ASSISTANT)
+
+fun user(context: String): Message = context.message(Role.USER)
+
+fun youAre(act: String, talkingTo: String, role: Role = Role.ASSISTANT): Message =
+  "You are a $act talking with a $talkingTo".message(role)
+
+class StepsMessageBuilder : PromptBuilder() {
+  override fun preprocess(elements: List<Message>): List<Message> =
+    elements.mapIndexed { ix, elt -> "${ix + 1} - ${elt.content}".message(elt.role) }
 }
 
-fun steps(inside: PromptBuilder.() -> Unit): Prompt = StepsBuilder().apply { inside() }.build()
+fun steps(inside: PromptBuilder.() -> Unit): List<Message> =
+  StepsMessageBuilder().apply { inside() }.build()
 
-fun writeSequenceOf(content: String): Prompt =
+fun writeSequenceOf(
+  content: String,
+  prefix: String = "Step",
+  role: Role = Role.ASSISTANT
+): Message =
   """
             Write a sequence of $content in the following format:
-            Step 1 - ...
-            Step 2 - ...
+            $prefix 1 - ...
+            $prefix 2 - ...
             ...
-            Step N - ...
-        """
-    .trimIndent()
-    .prompt()
-
-fun writeListOf(content: String): Prompt =
-  """
-            Write a list of $content in the following format:
-            1. ...
-            2. ...
-            n. ...
+            $prefix N - ...
         """
     .trimIndent()
-    .prompt()
+    .message(role)
 
-fun code(code: String, delimiter: Delimiter?, name: String? = null): Prompt =
+fun code(
+  code: String,
+  delimiter: Delimiter?,
+  name: String? = null,
+  role: Role = Role.ASSISTANT
+): Message =
   """
     ${name ?: "" }
     ${delimiter?.start() ?: ""}
     $code
     ${delimiter?.end() ?: ""}
     """
     .trimIndent()
-    .prompt()
+    .message(role)
 
 enum class Delimiter {
   ThreeBackticks,