babyagi_scala

main.scala

import openai.{OpenAiClient, OpenAiMockClient, OpenAiRestClient}
import pinecone.{PineconeClient, PineconeMockClient, PineconeRestClient}
import upickle.default._

import scala.annotation.tailrec
import scala.collection.immutable

case class Task(id: Int, name: String)
case class TaskContext(objective: String, tasks: Vector[Task])

object TaskContext {
  def apply(objective: String, task: Task): TaskContext = {
    TaskContext(objective, Vector(task))
  }
}

case class ContextAgent(tasks: Vector[String])

object Main {
  private val tableName = "test-table"

  private val openAiClient: OpenAiClient = OpenAiRestClient()
  private val pineconeClient: PineconeClient = PineconeRestClient()


  def getAdaEmbedding(text: String): Vector[Double] =
    val input = Vector(text.replaceAll("\n", " "))
    val response = openAiClient.createEmbedding(input, "text-embedding-ada-002")
    response.embedding

  def executionAgent(objective: String, task: String): String =
    val context = contextAgent(objective, tableName, 5)
//    println("\n*******RELEVANT CONTEXT******\n")
//    println(context.tasks.mkString(", "))
    val response = openAiClient.createCompletion(
      engine = "text-davinci-003",
      prompt = s"You are an AI who performs one task based on the following objective: ${objective}. Your task: ${task}\nResponse:",
      temperature = 0.7,
      maxTokens = 2000,
    )
    response.choices.head.trim

  def contextAgent(query: String, table: String, n: Int): ContextAgent =
    val queryEmbedding: Array[Double] = openAiClient.getAdaEmbedding(query)
    val result = pineconeClient.query(table, queryEmbedding, n, true)
    ContextAgent(result)

  def taskCreationAgent(objective: String, result: Map[String, String], taskDescription: String, taskList: Vector[Task]): Vector[Task] =
    val resultJson = write(result)
    val taskListJson = write(taskList.map(_.name))
    val prompt = s"You are an task creation AI that uses the result of an execution agent to create new tasks with the following objective: $objective, The last completed task has the result: $resultJson. This result was based on this task description: ${taskDescription}. These are incomplete tasks: $taskListJson. Based on the result, create new tasks to be completed by the AI system that do not overlap with incomplete tasks. Return the tasks as an array."
    val response = openAiClient.createCompletion("text-davinci-003", prompt, 0.5, 100)
    val newTasks = response.choices.head.strip().split('\n')
    newTasks.foldLeft(Vector[Task]()) { (acc, v) =>
      val newId: Int = (acc ++ taskList).maxByOption(_.id).map(_.id).getOrElse(0) + 1
      acc :+ Task(newId, v)
    }

  def prioritizationAgent(context: TaskContext, tasks: Vector[Task], lastTaskId: Int): TaskContext =
    val taskListJson = write(tasks.map(_.name))
    val nextTaskId = lastTaskId + 1
    val prompt: String =
      s"""You are an task prioritization AI tasked with cleaning the formatting of and reprioritizing the following tasks: $taskListJson. Consider the ultimate objective of your team:${context.objective}. Do not remove any tasks. Return the result as a numbered list, like:
        #. First task
        #. Second task
        Start the task list with number $nextTaskId."""
    val response = openAiClient.createCompletion("text-davinci-003", prompt, 0.5, 1000)
    val newTasks = response.choices.head.strip().split('\n').flatMap(line =>
      val taskParts = line.strip().split("\\.")
      if taskParts.length == 2 then
        val taskId = taskParts(0).strip()
        val taskName = taskParts(1).strip()
        Some(Task(taskId.toInt, taskName))
      else
        None
    ).toVector

    context.copy(tasks = newTasks)

  @tailrec
  def run(context: TaskContext): Unit = {
    context.tasks.headOption match
      case Some(task) =>
        val remainTasks = context.tasks.tail
        // Print the task list
        println("\u001b[95m\u001b[1m" + "\n*****TASK LIST*****\n" + "\u001b[0m\u001b[0m")
        context.tasks.foreach(t => println(s"${t.id}: ${t.name}"))

        // Step 1: Pull the first task
        println("\u001b[92m\u001b[1m" + "\n*****NEXT TASK*****\n" + "\u001b[0m\u001b[0m")
        println(s"${task.id}: ${task.name}")

        // Send to execution function to complete the task based on the context
        val result = executionAgent(context.objective, task.name)
        val thisTaskId = task.id
        println("\u001b[93m\u001b[1m" + "\n*****TASK RESULT*****\n" + "\u001b[0m\u001b[0m")
        println(result)

        // Step 2: Enrich result and store in Pinecone
        val enrichedResult = Map("data" -> result) // This is where you should enrich the result if needed
        val resultId = s"result_${task.id}"
        val vector = enrichedResult("data") // extract the actual result from the dictionary
        val adaEmbedding = getAdaEmbedding(vector)
        pineconeClient.upsert(resultId, adaEmbedding, task.name, result)

        // Step 3: Create new tasks and reprioritize task list
        val newTasks = taskCreationAgent(context.objective, enrichedResult, task.name, remainTasks)
        val nextContext = prioritizationAgent(context, remainTasks ++ newTasks, thisTaskId)
        run(nextContext)
      case _ =>
        println("Done.")
  }

  def main(args: Array[String]): Unit = {
    val context = TaskContext("Solve world hunger.", Task(1, "Develop a task list."))
    println("\u001b[96m\u001b[1m" + "\n*****OBJECTIVE*****\n" + "\u001b[0m\u001b[0m")
    println(context.objective)
    run(context)
  }
}

OpenAiRestClient.scala

package openai

import configure.Configure
import upickle.default._

case class CreateCompletionResponse(choices: Vector[String])
case class CreateEmbeddingResponse(embedding: Vector[Double])

trait OpenAiClient {
  def getAdaEmbedding(text: String): Array[Double]
  def createCompletion(engine: String, prompt: String, temperature: Double, maxTokens: Int, topP: Int = 1, frequencyPenalty: Double = 0.0, presencePenalty: Double = 0.0): CreateCompletionResponse
  def createEmbedding(input: Vector[String], model: String): CreateEmbeddingResponse
}

class OpenAiRestClient extends OpenAiClient {
  private def requestToOpenAi(url: String, requestData: String): String = {
    val headers = Map(
      "Content-Type" -> "application/json",
      "Authorization" -> s"Bearer ${Configure.openaiApiKey}"
    )

    val response = requests.post(url, headers = headers, data = requestData, connectTimeout = 100000, readTimeout = 100000)
    val responseText = response.text()

    responseText
  }

  def getAdaEmbedding(text: String): Array[Double] = {
    val replacedText = text.replace("\n", " ")

    val params = Map(
      "input" -> replacedText,
      "model" -> "text-embedding-ada-002"
    )

    val requestData = write(params)
    val responseText = requestToOpenAi("https://api.openai.com/v1/embeddings", requestData)

    val json = ujson.read(responseText)

    val embedding = read[Array[Double]](json("data")(0)("embedding"))
    embedding
  }

  def createCompletion(engine: String, prompt: String, temperature: Double, maxTokens: Int, topP: Int, frequencyPenalty: Double = 0.0, presencePenalty: Double): CreateCompletionResponse = {
    case class CreateCompletionRequest(
      model: String,
      prompt: String,
      temperature: Double,
      max_tokens: Int,
      top_p: Double,
      frequency_penalty: Double,
      presence_penalty: Double
    ) derives ReadWriter

    val params = CreateCompletionRequest(
      engine,
      prompt,
      temperature,
      maxTokens,
      topP,
      frequencyPenalty,
      presencePenalty
    )

    val requestData = write(params)
    val responseText = requestToOpenAi("https://api.openai.com/v1/completions", requestData)

    case class ChoiceData(text: String) derives ReadWriter
    case class JsonData(choices: Seq[ChoiceData]) derives ReadWriter

    val jsonData = read[JsonData](responseText)
    CreateCompletionResponse(jsonData.choices.map(_.text.trim).toVector)
  }

  def createEmbedding(input: Vector[String], model: String): CreateEmbeddingResponse = {
    case class CreateEmbeddingRequest(
      input: Vector[String],
      model: String
    ) derives ReadWriter

    val params = CreateEmbeddingRequest(
      input,
      model
    )

    val requestData = write(params)
    val responseText = requestToOpenAi("https://api.openai.com/v1/embeddings", requestData)

    case class DataData(embedding: Vector[Double]) derives ReadWriter
    case class JsonData(data: Seq[DataData]) derives ReadWriter

    val jsonData = read[JsonData](responseText)
    CreateEmbeddingResponse(jsonData.data.head.embedding)
  }
}

PineconeRestClient.scala

package pinecone

import configure.Configure
import upickle.default._

trait PineconeClient {
  def query(table: String, embedding: Array[Double], topK: Int, includeMetadata: Boolean): Vector[String]
  def upsert(id: String, embedding: Vector[Double], taskName: String, result: String): Unit
}


class PineconeRestClient extends PineconeClient {
  private def requestToPinecone(url: String, requestData: String): String = {
    val headers = Map(
      "Accept" -> "application/json",
      "Content-Type" -> "application/json",
      "Api-Key" -> s"${Configure.pineconeApiKey}"
    )

    val response = requests.post(url, headers = headers, data = requestData)
    val responseText = response.text()

    responseText
  }

  def query(table: String, embedding: Array[Double], topK: Int, includeMetadata: Boolean): Vector[String] = {
    case class RequestJsonObject(
      includeValues: Boolean,
      includeMetadata: Boolean,
      vector: Array[Double],
      topK: Int
    ) derives ReadWriter

    val params = RequestJsonObject(
      false,
      includeMetadata,
      embedding,
      topK,
    )
    val requestData = write(params)
    val responseText: String = requestToPinecone(s"${Configure.pineconeApiUrl}/query", requestData)

    case class Metadata(task: String) derives ReadWriter
    case class MatchData(id: String, score: Double, metadata: Metadata) derives ReadWriter
    case class JsonData(matches: Seq[MatchData]) derives ReadWriter

    val jsonData = read[JsonData](responseText)
    val matches = jsonData.matches

    val sortedMatches = matches.sortBy(-_.score)
    val result = sortedMatches.map(_.metadata.task).toVector

    result
  }

  def upsert(id: String, embedding: Vector[Double], taskName: String, result: String): Unit = {
    case class RequestJsonObject(
      id: String,
      values: Vector[Double],
      metadata: Map[String, String]
    ) derives ReadWriter

    val params = RequestJsonObject(
      id,
      embedding,
      Map(
        "task" -> taskName,
        "result" -> result
      ),
    )
    val requestData = write(Map("vectors" -> Seq(params)))
    requestToPinecone(s"${Configure.pineconeApiUrl}/vectors/upsert", requestData)
  }
}