ais: improve response schema + simplify the output of the result tool

getkyo · Nov 27, 2023 · 4e51009 · 4e51009
1 parent cf9b1cb
commit 4e51009
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 10 deletions.
diff --git a/kyo-llm-macros/shared/src/main/scala-2/kyo/llm/ValueSchema.scala b/kyo-llm-macros/shared/src/main/scala-2/kyo/llm/ValueSchema.scala
@@ -4,8 +4,6 @@ import zio.schema._
 import scala.language.experimental.macros
 import scala.reflect.macros.blackbox
 
-case class Value[T](value: T)
-
 case class ValueSchema[T](get: Schema[Value[T]])
 
 object ValueSchema {

diff --git a/kyo-llm-macros/shared/src/main/scala-3/kyo/llm/ValueSchema.scala b/kyo-llm-macros/shared/src/main/scala-3/kyo/llm/ValueSchema.scala
@@ -2,8 +2,6 @@ package kyo.llm
 
 import zio.schema._
 
-case class Value[T](value: T)
-
 case class ValueSchema[T](get: Schema[Value[T]])
 
 object ValueSchema {

diff --git a/kyo-llm-macros/shared/src/main/scala/kyo/llm/Value.scala b/kyo-llm-macros/shared/src/main/scala/kyo/llm/Value.scala
@@ -0,0 +1,17 @@
+package kyo.llm
+
+import scala.annotation.StaticAnnotation
+import zio.schema._
+
+final case class desc(value: String) extends StaticAnnotation
+
+case class Value[T](
+    @desc("Please **generate compact json**.")
+    willIGenerateCompactJson: Boolean,
+    @desc("Result is wrapped into a `value` field.")
+    value: T
+)
+
+object Value {
+  def apply[T](v: T): Value[T] = new Value(true, v)
+}
diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/ais.scala b/kyo-llm/shared/src/main/scala/kyo/llm/ais.scala
@@ -16,7 +16,6 @@ import kyo.tries._
 import zio.schema.codec.JsonCodec
 
 import java.lang.ref.WeakReference
-import scala.annotation.StaticAnnotation
 import scala.util.Failure
 import scala.util.Success
 import scala.util.control.NoStackTrace
@@ -28,7 +27,8 @@ object ais {
 
   type AIs >: AIs.Effects <: AIs.Effects
 
-  final case class desc(value: String) extends StaticAnnotation
+  type desc = kyo.llm.desc
+  val desc = kyo.llm.desc
 
   implicit class PromptInterpolator(val sc: StringContext) extends AnyVal {
     def p(args: Any*): String =

diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/tools.scala b/kyo-llm/shared/src/main/scala/kyo/llm/tools.scala
@@ -68,7 +68,7 @@ package object tools {
     )(f: (AI, T) => U > AIs)(implicit t: ValueSchema[T], u: ValueSchema[U]): Tool[T, U] =
       Tool(
           name,
-          description + " **Note how the input and output are wrapped into a `value` field**",
+          description,
           JsonSchema(t.get),
           JsonCodec.jsonDecoder(t.get),
           JsonCodec.jsonEncoder(u.get),
@@ -83,15 +83,15 @@ package object tools {
 
     private[kyo] def resultTool[T](implicit
         t: ValueSchema[T]
-    ): (Tool[T, T], Option[T] > AIs) > AIs = {
+    ): (Tool[T, String], Option[T] > AIs) > AIs = {
       Atomics.initRef(Option.empty[T]).map { ref =>
         val tool =
-          init[T, T](
+          init[T, String](
               "resultTool",
               "call this function with the result",
               (_: T) => "Generating result."
           ) { (_, v) =>
-            ref.set(Some(v)).andThen(v)
+            ref.set(Some(v)).andThen("Result processed.")
           }
         (tool, ref.get)
       }