From ac94fb173bab3c297594a3408e118c4f6b8d4f99 Mon Sep 17 00:00:00 2001 From: Flavio Brasil Date: Sat, 30 Dec 2023 09:31:18 -0800 Subject: [PATCH] thoughts check + observe wip --- .../src/main/scala/kyo/stats/attributes.scala | 5 +- .../scala/kyoTest/stats/AttributesTest.scala | 28 +++--- .../scala/kyoTest/stats/CounterTest.scala | 6 +- .../scala/kyoTest/stats/HistogramTest.scala | 6 +- .../src/main/scala/kyo/llm/agents.scala | 10 +-- .../src/main/scala/kyo/llm/completions.scala | 2 +- .../main/scala/kyo/llm/thoughts/Check.scala | 85 +++++++++++++------ .../main/scala/kyo/llm/thoughts/Observe.scala | 56 ++++++++++++ .../main/scala/kyo/llm/thoughts/Thought.scala | 12 +-- .../stats/otel/OTelAttributesTest.scala | 18 ++-- .../kyoTest/stats/otel/OTelReceiverTest.scala | 4 +- 11 files changed, 161 insertions(+), 71 deletions(-) create mode 100644 kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Observe.scala diff --git a/kyo-core/shared/src/main/scala/kyo/stats/attributes.scala b/kyo-core/shared/src/main/scala/kyo/stats/attributes.scala index 8cc7ea52c..53c06f5a9 100644 --- a/kyo-core/shared/src/main/scala/kyo/stats/attributes.scala +++ b/kyo-core/shared/src/main/scala/kyo/stats/attributes.scala @@ -1,16 +1,19 @@ package kyo.stats import scala.annotation.implicitNotFound +import kyo.stats.Attributes.AsAttribute case class Attributes(get: List[Attributes.Attribute]) extends AnyVal { def add(a: Attributes): Attributes = Attributes(get ++ a.get) + def add[T](name: String, value: T)(implicit a: AsAttribute[T]): Attributes = + add(Attributes.add(name, value)) } object Attributes { val empty: Attributes = Attributes(Nil) - def of[T](name: String, value: T)(implicit a: AsAttribute[T]) = + def add[T](name: String, value: T)(implicit a: AsAttribute[T]) = Attributes(a.f(name, value) :: Nil) def all(l: List[Attributes]): Attributes = diff --git a/kyo-core/shared/src/test/scala/kyoTest/stats/AttributesTest.scala b/kyo-core/shared/src/test/scala/kyoTest/stats/AttributesTest.scala index e2b58ecd5..09880a9cd 100644 --- a/kyo-core/shared/src/test/scala/kyoTest/stats/AttributesTest.scala +++ b/kyo-core/shared/src/test/scala/kyoTest/stats/AttributesTest.scala @@ -13,20 +13,20 @@ class AttributesTest extends KyoTest { } "one" in { - val attr = Attributes.of("test", true) + val attr = Attributes.add("test", true) assert(attr.get.size == 1) assert(attr.get.head.isInstanceOf[Attribute.BooleanAttribute]) } "add" in { - val attr1 = Attributes.of("test1", true) - val attr2 = Attributes.of("test2", 123) + val attr1 = Attributes.add("test1", true) + val attr2 = Attributes.add("test2", 123) val combined = attr1.add(attr2) assert(combined.get.size == 2) } "all" in { - val attrList = List(Attributes.of("test1", true), Attributes.of("test2", 123.45)) + val attrList = List(Attributes.add("test1", true), Attributes.add("test2", 123.45)) val combined = Attributes.all(attrList) assert(combined.get.size == attrList.size) } @@ -34,58 +34,58 @@ class AttributesTest extends KyoTest { "primitives" - { "boolean" in { - val booleanAttr = Attributes.of("bool", true) + val booleanAttr = Attributes.add("bool", true) assert(booleanAttr.get.head.isInstanceOf[Attribute.BooleanAttribute]) } "int" in { - val booleanAttr = Attributes.of("int", 1) + val booleanAttr = Attributes.add("int", 1) assert(booleanAttr.get.head.isInstanceOf[Attribute.LongAttribute]) } "double" in { - val doubleAttr = Attributes.of("double", 123.45) + val doubleAttr = Attributes.add("double", 123.45) assert(doubleAttr.get.head.isInstanceOf[Attribute.DoubleAttribute]) } "long" in { - val longAttr = Attributes.of("long", 123L) + val longAttr = Attributes.add("long", 123L) assert(longAttr.get.head.isInstanceOf[Attribute.LongAttribute]) } "string" in { - val stringAttr = Attributes.of("string", "value") + val stringAttr = Attributes.add("string", "value") assert(stringAttr.get.head.isInstanceOf[Attribute.StringAttribute]) } } "lists" - { "boolean list" in { - val boolListAttr = Attributes.of("boolList", List(true, false, true)) + val boolListAttr = Attributes.add("boolList", List(true, false, true)) assert(boolListAttr.get.head.isInstanceOf[Attribute.BooleanListAttribute]) assert(boolListAttr.get.head.asInstanceOf[Attribute.BooleanListAttribute].value.size == 3) } "integer list" in { - val intListAttr = Attributes.of("intList", List(1, 2, 3)) + val intListAttr = Attributes.add("intList", List(1, 2, 3)) assert(intListAttr.get.head.isInstanceOf[Attribute.LongListAttribute]) assert(intListAttr.get.head.asInstanceOf[Attribute.LongListAttribute].value.size == 3) } "double list" in { - val doubleListAttr = Attributes.of("doubleList", List(1.1, 2.2, 3.3)) + val doubleListAttr = Attributes.add("doubleList", List(1.1, 2.2, 3.3)) assert(doubleListAttr.get.head.isInstanceOf[Attribute.DoubleListAttribute]) assert(doubleListAttr.get.head.asInstanceOf[Attribute.DoubleListAttribute].value.size == 3) } "long list" in { - val longListAttr = Attributes.of("longList", List(100L, 200L, 300L)) + val longListAttr = Attributes.add("longList", List(100L, 200L, 300L)) assert(longListAttr.get.head.isInstanceOf[Attribute.LongListAttribute]) assert(longListAttr.get.head.asInstanceOf[Attribute.LongListAttribute].value.size == 3) } "string list" in { - val stringListAttr = Attributes.of("stringList", List("a", "b", "c")) + val stringListAttr = Attributes.add("stringList", List("a", "b", "c")) assert(stringListAttr.get.head.isInstanceOf[Attribute.StringListAttribute]) assert(stringListAttr.get.head.asInstanceOf[Attribute.StringListAttribute].value.size == 3) } diff --git a/kyo-core/shared/src/test/scala/kyoTest/stats/CounterTest.scala b/kyo-core/shared/src/test/scala/kyoTest/stats/CounterTest.scala index dd2d9a703..53a2f9a6a 100644 --- a/kyo-core/shared/src/test/scala/kyoTest/stats/CounterTest.scala +++ b/kyo-core/shared/src/test/scala/kyoTest/stats/CounterTest.scala @@ -10,7 +10,7 @@ class CounterTest extends KyoTest { for { _ <- Counter.noop.inc _ <- Counter.noop.add(1) - _ <- Counter.noop.add(1, Attributes.of("test", 1)) + _ <- Counter.noop.add(1, Attributes.add("test", 1)) } yield succeed } @@ -20,7 +20,7 @@ class CounterTest extends KyoTest { for { _ <- counter.inc _ <- counter.add(1) - _ <- counter.add(1, Attributes.of("test", 1)) + _ <- counter.add(1, Attributes.add("test", 1)) } yield assert(unsafe.curr == 3) } @@ -39,7 +39,7 @@ class CounterTest extends KyoTest { for { _ <- counter.inc _ <- counter.add(1) - _ <- counter.add(1, Attributes.of("test", 1)) + _ <- counter.add(1, Attributes.add("test", 1)) } yield { assert(unsafe1.curr == 3 && unsafe2.curr == 3) } diff --git a/kyo-core/shared/src/test/scala/kyoTest/stats/HistogramTest.scala b/kyo-core/shared/src/test/scala/kyoTest/stats/HistogramTest.scala index 4d8a4a6f9..ba0912937 100644 --- a/kyo-core/shared/src/test/scala/kyoTest/stats/HistogramTest.scala +++ b/kyo-core/shared/src/test/scala/kyoTest/stats/HistogramTest.scala @@ -9,7 +9,7 @@ class HistogramTest extends KyoTest { "noop" in run { for { _ <- Histogram.noop.observe(1.0) - _ <- Histogram.noop.observe(1.0, Attributes.of("test", 1)) + _ <- Histogram.noop.observe(1.0, Attributes.add("test", 1)) } yield succeed } @@ -18,7 +18,7 @@ class HistogramTest extends KyoTest { val histogram = Histogram(unsafe) for { _ <- histogram.observe(1.0) - _ <- histogram.observe(1.0, Attributes.of("test", 1)) + _ <- histogram.observe(1.0, Attributes.add("test", 1)) } yield assert(unsafe.observations == 2) } @@ -36,7 +36,7 @@ class HistogramTest extends KyoTest { val histogram = Histogram.all(List(Histogram(unsafe1), Histogram(unsafe2))) for { _ <- histogram.observe(1.0) - _ <- histogram.observe(1.0, Attributes.of("test", 1)) + _ <- histogram.observe(1.0, Attributes.add("test", 1)) } yield { assert(unsafe1.observations == 2 && unsafe2.observations == 2) } diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/agents.scala b/kyo-llm/shared/src/main/scala/kyo/llm/agents.scala index 89991f312..b3a80ecda 100644 --- a/kyo-llm/shared/src/main/scala/kyo/llm/agents.scala +++ b/kyo-llm/shared/src/main/scala/kyo/llm/agents.scala @@ -38,9 +38,9 @@ package object agents { val output: Json[Out] ) - def info: Info + val info: Info - def thoughts: List[Thought.Info] = Nil + val thoughts: List[Thought.Info] = Nil private val local = Locals.init(Option.empty[AI]) @@ -57,7 +57,7 @@ package object agents { case None => AIs.init } - private[kyo] def request: Schema = { + private[kyo] val schema: Schema = { def schema[T](name: String, l: List[Thought.Info]): ZSchema[T] = { val fields = l.map { t => import zio.schema.Schema._ @@ -68,7 +68,7 @@ package object agents { Validation.succeed, identity, (_, _) => ListMap.empty - ) + ) } ZSchema.record(TypeId.fromTypeName(name), FieldSet(fields: _*)).asInstanceOf[ZSchema[T]] } @@ -153,7 +153,7 @@ package object agents { "Call this agent with the result." ) - override def thoughts: List[Thought.Info] = + override val thoughts: List[Thought.Info] = _thoughts def run(input: T) = diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala b/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala index 1bf55ecec..735d741cf 100644 --- a/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala +++ b/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala @@ -173,7 +173,7 @@ object completions { ToolDef(FunctionDef( p.info.description, p.info.name, - p.request + p.schema )) ).toList) Request( diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Check.scala b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Check.scala index fbe5db9e5..59067e05a 100644 --- a/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Check.scala +++ b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Check.scala @@ -9,32 +9,45 @@ import kyo.ios.IOs object Check { - private val stats = Thought.stats.scope("checks") + private val stats = Thought.stats.scope("check") private val success = stats.initCounter("success") private val failure = stats.initCounter("failure") - case class CheckFailed(path: List[String], invariant: String, analysis: String) + case class CheckFailed(ai: AI, thought: Thought, invariant: String, analysis: String) extends RuntimeException - private def observe(path: List[String], result: Boolean) = { + private def observe(parent: Thought, field: String, result: Boolean) = { val c = if (result) success else failure - c.attributes(Attributes.of("thought", path.last)).inc + c.attributes( + Attributes + .add("thought", parent.name) + .add("field", field) + ).inc } - private def warn(ai: AI, path: List[String], invariant: String): Unit < AIs = + private def warn( + ai: AI, + parent: Thought, + field: String, + invariant: String, + analysis: String = "Plase reason about the failure and fix any mistakes.", + repair: Option[Repair] = None + ): Unit < AIs = ai.systemMessage( p""" Thought Invariant Failure ========================= + Thought: ${parent.name} + Field: $field Description: $invariant - Path: ${path.map(v => s"`$v`").mkString(".")} - Plase analyze and fix any mistakes. + Analysis: $analysis + ${repair.map(pprint(_).plainText).map("Repair: " + _).getOrElse("")} """ ) case class Info(result: Boolean) extends Thought { - override def eval(path: List[String], ai: AI) = - observe(path, result) + override def eval(parent: Thought, field: String, ai: AI) = + observe(parent, field, result) } object Info { @@ -46,9 +59,9 @@ object Check { `Invariant check description`: Invarant, `Invariant holds`: Boolean ) extends Thought { - override def eval(path: List[String], ai: AI) = - observe(path, `Invariant holds`).andThen { - warn(ai, path, `Invariant check description`) + override def eval(parent: Thought, field: String, ai: AI) = + observe(parent, field, `Invariant holds`).andThen { + warn(ai, parent, field, `Invariant check description`) } } @@ -57,10 +70,21 @@ object Check { `Invariant check analysis`: String, `Invariant holds`: Boolean ) extends Thought { - override def eval(path: List[String], ai: AI) = - observe(path, `Invariant holds`).andThen { - warn(ai, path, `Invariant check description`).andThen { - IOs.fail(CheckFailed(path, `Invariant check description`, `Invariant check analysis`)) + override def eval(parent: Thought, field: String, ai: AI) = + observe(parent, field, `Invariant holds`).andThen { + warn( + ai, + parent, + field, + `Invariant check description`, + `Invariant check analysis` + ).andThen { + IOs.fail(CheckFailed( + ai, + parent, + `Invariant check description`, + `Invariant check analysis` + )) } } } @@ -70,21 +94,26 @@ object Check { `Invariant check analysis`: String, `Invariant holds`: Boolean ) extends Thought { - override def eval(path: List[String], ai: AI) = - observe(path, `Invariant holds`).andThen { + override def eval(parent: Thought, field: String, ai: AI) = + observe(parent, field, `Invariant holds`).andThen { AIs.ephemeral { - warn(ai, path, `Invariant check description`).andThen { - ai.gen[Repair]("Provide a repair for the failed thought invariant.") + warn( + ai, + parent, + field, + `Invariant check description`, + `Invariant check analysis` + ).andThen { + ai.gen[Repair]("Provide a repair for the last failed thought invariant.") } }.map { repair => - ai.systemMessage( - p""" - Thought Invariant Repair - ======================== - Description: ${`Invariant check description`} - Path: ${path.map(v => s"`$v`").mkString(".")} - Inferred Repair: ${pprint(repair)} - """ + warn( + ai, + parent, + field, + `Invariant check description`, + `Invariant check analysis`, + Some(repair) ) } } diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Observe.scala b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Observe.scala new file mode 100644 index 000000000..25e0ca534 --- /dev/null +++ b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Observe.scala @@ -0,0 +1,56 @@ +package kyo.llm.thoughts + +import kyo.llm.ais._ +import zio.schema.{Schema => ZSchema} +import kyo._ +import kyo.ios._ +import kyo.stats._ + +object Observe { + + private val stats = Thought.stats.scope("observe") + private val count = stats.initCounter("count") + private val sum = stats.initCounter("sum") + private val hist = stats.initHistogram("distribution") + + private def add(counter: Counter, parent: Thought, field: String, v: Long) = + counter.attributes( + Attributes + .add("thought", parent.name) + .add("field", field) + ).add(v) + + case class Count[T](v: T) extends Thought { + override def eval(parent: Thought, field: String, ai: AI) = + add(count, parent, field, 1) + } + + object Count { + implicit def schema[T](implicit s: ZSchema[T]): ZSchema[Count[T]] = + s.transform(Count(_), _.v) + } + + case class Sum(v: Int) extends Thought { + override def eval(parent: Thought, field: String, ai: AI) = + add(sum, parent, field, v) + } + + object Sum { + implicit val schema: ZSchema[Sum] = + ZSchema.primitive[Int].transform(Sum(_), _.v) + } + + case class Distribution(v: Double) extends Thought { + override def eval(parent: Thought, field: String, ai: AI) = + hist.attributes( + Attributes + .add("thought", parent.name) + .add("field", field) + ).observe(v) + } + + object Distribution { + implicit val schema: ZSchema[Distribution] = + ZSchema.primitive[Double].transform(Distribution(_), _.v) + } +} diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Thought.scala b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Thought.scala index db4183c38..d9f06c689 100644 --- a/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Thought.scala +++ b/kyo-llm/shared/src/main/scala/kyo/llm/thoughts/Thought.scala @@ -10,10 +10,11 @@ import kyo.stats.Stats abstract class Thought { self: Product => - final def handle(path: List[String], ai: AI): Unit < AIs = { - val npath = productPrefix :: path + def name = productPrefix + + final def handle(parent: Thought, field: String, ai: AI): Unit < AIs = { val thoughts = - eval(npath, ai).andThen { + eval(parent, field, ai).andThen { (0 until productArity).flatMap { idx => productElement(idx) match { case v: Thought => @@ -23,10 +24,11 @@ abstract class Thought { } } } - AIs.parallelTraverse(thoughts)((f, t) => t.eval(f :: npath, ai)).unit + AIs.parallelTraverse(thoughts)((f, t) => t.eval(self, f, ai)).unit } - def eval(path: List[String], ai: AI): Unit < AIs = () + def eval(parent: Thought, field: String, ai: AI): Unit < AIs = () + } object Thought { diff --git a/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelAttributesTest.scala b/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelAttributesTest.scala index 2b2de7724..2e1b370e3 100644 --- a/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelAttributesTest.scala +++ b/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelAttributesTest.scala @@ -10,15 +10,15 @@ class OTelAttributesTest extends KyoTest { "test" in { val kyoAttrs = Attributes.all(List( - Attributes.of("boolAttr", true), - Attributes.of("doubleAttr", 2.0), - Attributes.of("intAttr", 42), - Attributes.of("longAttr", 123L), - Attributes.of("stringAttr", "test"), - Attributes.of("boolListAttr", List(true, false)), - Attributes.of("doubleListAttr", List(1.1, 2.2)), - Attributes.of("longListAttr", List(100L, 200L)), - Attributes.of("stringListAttr", List("a", "b")) + Attributes.add("boolAttr", true), + Attributes.add("doubleAttr", 2.0), + Attributes.add("intAttr", 42), + Attributes.add("longAttr", 123L), + Attributes.add("stringAttr", "test"), + Attributes.add("boolListAttr", List(true, false)), + Attributes.add("doubleListAttr", List(1.1, 2.2)), + Attributes.add("longListAttr", List(100L, 200L)), + Attributes.add("stringListAttr", List("a", "b")) )) val oTelAttrs = OTelAttributes(kyoAttrs) diff --git a/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelReceiverTest.scala b/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelReceiverTest.scala index fadc897a4..eee3c6ae4 100644 --- a/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelReceiverTest.scala +++ b/kyo-stats-otel/src/test/scala/kyoTest/stats/otel/OTelReceiverTest.scala @@ -26,9 +26,9 @@ class OTelReceiverTest extends KyoTest { for { _ <- counter.inc _ <- counter.add(1) - _ <- counter.add(2, Attributes.of("test", 3)) + _ <- counter.add(2, Attributes.add("test", 3)) _ <- histogram.observe(42d) - _ <- histogram.observe(24d, Attributes.of("test", 3)) + _ <- histogram.observe(24d, Attributes.add("test", 3)) } yield succeed }