update observation

LLMAgentBuilder · Oct 5, 2024 · aae4e1c · aae4e1c
1 parent 99085ce
commit aae4e1c
Show file tree

Hide file tree

Showing 20 changed files with 191 additions and 35 deletions.
diff --git a/...er-executor/src/main/kotlin/io/github/llmagentbuilder/planner/executor/LLMPlanExecutor.kt b/...er-executor/src/main/kotlin/io/github/llmagentbuilder/planner/executor/LLMPlanExecutor.kt
@@ -4,16 +4,31 @@ import io.github.llmagentbuilder.core.AgentFinish
 import io.github.llmagentbuilder.core.IntermediateAgentStep
 import io.github.llmagentbuilder.core.Planner
 import io.github.llmagentbuilder.core.executor.ActionPlanningResult
+import io.github.llmagentbuilder.core.observation.AgentPlanningObservationContext
+import io.github.llmagentbuilder.core.observation.AgentPlanningObservationDocumentation
+import io.github.llmagentbuilder.core.observation.DefaultAgentPlanningObservationConvention
 import io.github.llmagentbuilder.core.planner.OutputParser
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
 open class LLMPlanExecutor(
     private val chatClient: ChatClient,
     private val outputParser: OutputParser,
+    private val observationRegistry: ObservationRegistry? = null,
 ) : Planner {
     override fun plan(
         inputs: Map<String, Any>,
         intermediateSteps: List<IntermediateAgentStep>
+    ): ActionPlanningResult {
+        val action = { internalPlan(inputs, intermediateSteps) }
+        return observationRegistry?.let { registry ->
+            instrumentedPlan(inputs, action, registry)
+        } ?: action.invoke()
+    }
+
+    private fun internalPlan(
+        inputs: Map<String, Any>,
+        intermediateSteps: List<IntermediateAgentStep>
     ): ActionPlanningResult {
         val userInput =
             (inputs["input"] as? String)
@@ -36,6 +51,34 @@ open class LLMPlanExecutor(
         return ActionPlanningResult.fromParseResult(result)
     }
 
+    private fun instrumentedPlan(
+        input: Map<String, Any>,
+        action: () -> ActionPlanningResult,
+        registry: ObservationRegistry
+    ): ActionPlanningResult {
+        val observationContext =
+            AgentPlanningObservationContext(input)
+        val observation =
+            AgentPlanningObservationDocumentation.AGENT_PLANNING.observation(
+                null,
+                DefaultAgentPlanningObservationConvention(),
+                { observationContext },
+                registry
+            ).start()
+        return try {
+            observation.openScope().use {
+                val response = action.invoke()
+                observationContext.setResponse(response)
+                response
+            }
+        } catch (e: Exception) {
+            observation.error(e)
+            throw e
+        } finally {
+            observation.stop()
+        }
+    }
+
     private fun constructScratchpad(intermediateSteps: List<IntermediateAgentStep>): String {
         return intermediateSteps.joinToString(" ") {
             val (action, observation) = it

diff --git a/...eact-json/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPlanner.kt b/...eact-json/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPlanner.kt
@@ -1,7 +1,15 @@
 package io.github.llmagentbuilder.planner.reactjson
 
 import io.github.llmagentbuilder.planner.executor.LLMPlanExecutor
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
-class ReActJsonPlanner(chatClient: ChatClient) :
-    LLMPlanExecutor(chatClient, ReActJsonOutputParser.INSTANCE)
+class ReActJsonPlanner(
+    chatClient: ChatClient,
+    observationRegistry: ObservationRegistry? = null,
+) :
+    LLMPlanExecutor(
+        chatClient,
+        ReActJsonOutputParser.INSTANCE,
+        observationRegistry
+    )
diff --git a/...n/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPlannerProvider.kt b/...n/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPlannerProvider.kt
@@ -3,6 +3,7 @@ package io.github.llmagentbuilder.planner.reactjson
 import io.github.llmagentbuilder.core.MapToObject
 import io.github.llmagentbuilder.core.Planner
 import io.github.llmagentbuilder.core.PlannerProvider
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
 class ReActJsonPlannerProvider : PlannerProvider {
@@ -12,7 +13,8 @@ class ReActJsonPlannerProvider : PlannerProvider {
 
     override fun providePlanner(
         chatClientBuilder: ChatClient.Builder,
-        config: Map<String, Any?>?
+        config: Map<String, Any?>?,
+        observationRegistry: ObservationRegistry?,
     ): Planner? {
         val plannerConfig = MapToObject.toObject<ReActJsonPlannerConfig>(config)
         if (plannerConfig?.enabled == false) {
@@ -21,6 +23,6 @@ class ReActJsonPlannerProvider : PlannerProvider {
         val chatClient =
             chatClientBuilder.defaultAdvisors(ReActJsonPromptAdvisor())
                 .build()
-        return ReActJsonPlanner(chatClient)
+        return ReActJsonPlanner(chatClient, observationRegistry)
     }
 }
diff --git a/...son/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPromptAdvisor.kt b/...son/src/main/kotlin/io/github/llmagentbuilder/planner/reactjson/ReActJsonPromptAdvisor.kt
@@ -56,7 +56,7 @@ internal const val defaultUserTextTemplate = """
 
 class ReActJsonPromptAdvisor : CallAroundAdvisor {
     override fun getName(): String {
-        return "ReAct Json Planner - Prompt"
+        return javaClass.simpleName
     }
 
     override fun getOrder(): Int {

diff --git a/...nts/planner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPlanner.kt b/...nts/planner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPlanner.kt
@@ -1,7 +1,11 @@
 package io.github.llmagentbuilder.planner.react
 
 import io.github.llmagentbuilder.planner.executor.LLMPlanExecutor
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
-class ReActPlanner(chatClient: ChatClient) :
-    LLMPlanExecutor(chatClient, ReActOutputParser.INSTANCE)
+class ReActPlanner(
+    chatClient: ChatClient,
+    observationRegistry: ObservationRegistry? = null,
+) :
+    LLMPlanExecutor(chatClient, ReActOutputParser.INSTANCE, observationRegistry)
diff --git a/...ner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPlannerProvider.kt b/...ner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPlannerProvider.kt
@@ -3,6 +3,7 @@ package io.github.llmagentbuilder.planner.react
 import io.github.llmagentbuilder.core.MapToObject
 import io.github.llmagentbuilder.core.Planner
 import io.github.llmagentbuilder.core.PlannerProvider
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
 class ReActPlannerProvider : PlannerProvider {
@@ -12,14 +13,15 @@ class ReActPlannerProvider : PlannerProvider {
 
     override fun providePlanner(
         chatClientBuilder: ChatClient.Builder,
-        config: Map<String, Any?>?
+        config: Map<String, Any?>?,
+        observationRegistry: ObservationRegistry?,
     ): Planner? {
         val plannerConfig = MapToObject.toObject<ReActPlannerConfig>(config)
         if (plannerConfig?.enabled == false) {
             return null
         }
         val chatClient = chatClientBuilder.defaultAdvisors(ReActPromptAdvisor())
             .build()
-        return ReActPlanner(chatClient)
+        return ReActPlanner(chatClient, observationRegistry)
     }
 }
diff --git a/...anner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPromptAdvisor.kt b/...anner/react/src/main/kotlin/io/github/llmagentbuilder/planner/react/ReActPromptAdvisor.kt
@@ -38,7 +38,7 @@ internal const val defaultUserTextTemplate = """
 
 class ReActPromptAdvisor : CallAroundAdvisor {
     override fun getName(): String {
-        return "ReAct Planner - Prompt"
+        return javaClass.simpleName
     }
 
     override fun getOrder(): Int {

diff --git a/.../planner/simple/src/main/kotlin/io/github/llmagentbuilder/planner/simple/SimplePlanner.kt b/.../planner/simple/src/main/kotlin/io/github/llmagentbuilder/planner/simple/SimplePlanner.kt
@@ -1,8 +1,15 @@
 package io.github.llmagentbuilder.planner.simple
 
 import io.github.llmagentbuilder.planner.executor.LLMPlanExecutor
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
-class SimplePlanner(chatClient: ChatClient) :
-    LLMPlanExecutor(chatClient, SimpleOutputParser.INSTANCE) {
-}
+class SimplePlanner(
+    chatClient: ChatClient,
+    observationRegistry: ObservationRegistry? = null
+) :
+    LLMPlanExecutor(
+        chatClient,
+        SimpleOutputParser.INSTANCE,
+        observationRegistry
+    )
diff --git a/.../simple/src/main/kotlin/io/github/llmagentbuilder/planner/simple/SimplePlannerProvider.kt b/.../simple/src/main/kotlin/io/github/llmagentbuilder/planner/simple/SimplePlannerProvider.kt
@@ -3,6 +3,7 @@ package io.github.llmagentbuilder.planner.simple
 import io.github.llmagentbuilder.core.MapToObject
 import io.github.llmagentbuilder.core.Planner
 import io.github.llmagentbuilder.core.PlannerProvider
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
 class SimplePlannerProvider : PlannerProvider {
@@ -12,12 +13,13 @@ class SimplePlannerProvider : PlannerProvider {
 
     override fun providePlanner(
         chatClientBuilder: ChatClient.Builder,
-        config: Map<String, Any?>?
+        config: Map<String, Any?>?,
+        observationRegistry: ObservationRegistry?,
     ): Planner? {
         val plannerConfig = MapToObject.toObject<SimplePlannerConfig>(config)
         if (plannerConfig?.enabled == false) {
             return null
         }
-        return SimplePlanner(chatClientBuilder.build())
+        return SimplePlanner(chatClientBuilder.build(), observationRegistry)
     }
 }
diff --git a/...src/main/kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPlanner.kt b/...src/main/kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPlanner.kt
@@ -1,7 +1,15 @@
 package io.github.llmagentbuilder.planner.structuredchat
 
 import io.github.llmagentbuilder.planner.executor.LLMPlanExecutor
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
-class StructuredChatPlanner(chatClient: ChatClient) :
-    LLMPlanExecutor(chatClient, StructuredChatOutputParser.INSTANCE)
+class StructuredChatPlanner(
+    chatClient: ChatClient,
+    observationRegistry: ObservationRegistry? = null,
+) :
+    LLMPlanExecutor(
+        chatClient,
+        StructuredChatOutputParser.INSTANCE,
+        observationRegistry
+    )
diff --git a/.../kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPlannerProvider.kt b/.../kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPlannerProvider.kt
@@ -3,6 +3,7 @@ package io.github.llmagentbuilder.planner.structuredchat
 import io.github.llmagentbuilder.core.MapToObject
 import io.github.llmagentbuilder.core.Planner
 import io.github.llmagentbuilder.core.PlannerProvider
+import io.micrometer.observation.ObservationRegistry
 import org.springframework.ai.chat.client.ChatClient
 
 class StructuredChatPlannerProvider : PlannerProvider {
@@ -12,7 +13,8 @@ class StructuredChatPlannerProvider : PlannerProvider {
 
     override fun providePlanner(
         chatClientBuilder: ChatClient.Builder,
-        config: Map<String, Any?>?
+        config: Map<String, Any?>?,
+        observationRegistry: ObservationRegistry?,
     ): Planner? {
         val plannerConfig =
             MapToObject.toObject<StructuredChatPlannerConfig>(config)
@@ -22,6 +24,6 @@ class StructuredChatPlannerProvider : PlannerProvider {
         val chatClient =
             chatClientBuilder.defaultAdvisors(StructuredChatPromptAdvisor())
                 .build()
-        return StructuredChatPlanner(chatClient)
+        return StructuredChatPlanner(chatClient, observationRegistry)
     }
 }
diff --git a/...in/kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPromptAdvisor.kt b/...in/kotlin/io/github/llmagentbuilder/planner/structuredchat/StructuredChatPromptAdvisor.kt
@@ -55,7 +55,7 @@ internal const val defaultUserTextTemplate = """
 
 class StructuredChatPromptAdvisor : CallAroundAdvisor {
     override fun getName(): String {
-        return "StructuredChat Planner - Prompt"
+        return javaClass.simpleName
     }
 
     override fun getOrder(): Int {

diff --git a/...tlin/io/github/llmagentbuilder/agent/profile/systemmessage/SystemMessageProfileAdvisor.kt b/...tlin/io/github/llmagentbuilder/agent/profile/systemmessage/SystemMessageProfileAdvisor.kt
@@ -11,7 +11,7 @@ class SystemMessageProfileAdvisor(
     private val systemMessageParams: Map<String, Any>? = mapOf()
 ) : CallAroundAdvisor {
     override fun getName(): String {
-        return "Profile - System Message"
+        return javaClass.simpleName
     }
 
     override fun getOrder(): Int {

diff --git a/...der/agent/tool/AgentToolContextAdvisor.kt → ...uilder/agent/tool/AgentToolInfoAdvisor.kt b/...der/agent/tool/AgentToolContextAdvisor.kt → ...uilder/agent/tool/AgentToolInfoAdvisor.kt
@@ -10,10 +10,10 @@ import org.springframework.core.Ordered
 const val SYSTEM_PARAM_TOOL_NAMES = "tool_names"
 const val SYSTEM_PARAM_TOOLS = "tools"
 
-class AgentToolContextAdvisor(private val tools: Map<String, AgentTool<*, *>>) :
+class AgentToolInfoAdvisor(private val tools: Map<String, AgentTool<*, *>>) :
     CallAroundAdvisor {
     override fun getName(): String {
-        return "AgentTool"
+        return javaClass.simpleName
     }
 
     override fun getOrder(): Int {

diff --git a/bootstrap/src/main/kotlin/io/github/llmagentbuilder/bootstrap/AgentBootstrap.kt b/bootstrap/src/main/kotlin/io/github/llmagentbuilder/bootstrap/AgentBootstrap.kt
@@ -1,7 +1,7 @@
 package io.github.llmagentbuilder.bootstrap
 
 import io.github.llmagentbuilder.agent.profile.systemmessage.SystemMessageProfileAdvisor
-import io.github.llmagentbuilder.agent.tool.AgentToolContextAdvisor
+import io.github.llmagentbuilder.agent.tool.AgentToolInfoAdvisor
 import io.github.llmagentbuilder.core.*
 import io.github.llmagentbuilder.core.tool.AgentToolFunctionCallbackContext
 import io.github.llmagentbuilder.core.tool.AgentToolsProviderFactory
@@ -40,17 +40,20 @@ object AgentBootstrap {
         val agentToolsProvider =
             AgentToolsProviderFactory.create(agentConfig.tools ?: listOf())
         advisors.addLast(
-            AgentToolContextAdvisor(
+            AgentToolInfoAdvisor(
                 agentToolsProvider.get()
             )
         )
         advisors.addLast(SimpleLoggerAdvisor())
+        val observationEnabled = agentConfig.observation?.enabled == true
+        val observationRegistry =
+            if (observationEnabled) ObservationRegistry.create() else ObservationRegistry.NOOP
         val functionCallbackContext =
             AgentToolFunctionCallbackContext(
                 agentToolsProvider,
+                observationRegistry,
             )
-        val observationRegistry = ObservationRegistry.create()
-        if (agentConfig.observation?.enabled == true) {
+        if (observationEnabled) {
             OpenTelemetryPlugin().install(agentConfig, observationRegistry)
         }
         val llmConfigs = agentConfig.llm
@@ -77,7 +80,11 @@ object AgentBootstrap {
             .mapNotNull {
                 val config =
                     plannerConfigs?.get(it.configKey()) as? Map<String, Any?>
-                it.providePlanner(chatClientBuilder, config)
+                it.providePlanner(
+                    chatClientBuilder,
+                    config,
+                    observationRegistry,
+                )
             }
             .firstOrNull()?.also {
                 logger.info("Loaded Planner $it")

diff --git a/bootstrap/src/main/resources/agent.yaml b/bootstrap/src/main/resources/agent.yaml
@@ -12,11 +12,19 @@ planner:
   reActJson:
     enabled: true
 observation:
+  enabled: true
   tracing:
+    enabled: true
     exporter:
       endpoint: "https://api.honeycomb.io/v1/traces"
       headers:
         "x-honeycomb-team": "{{env.HONEYCOMB_API_KEY}}"
+  metrics:
+    enabled: true
+    exporter:
+      endpoint: "https://api.honeycomb.io/v1/metrics"
+      headers:
+        "x-honeycomb-team": "{{env.HONEYCOMB_API_KEY}}"
 tools:
   - id: writeLocalFile
     config:

diff --git a/core/src/main/kotlin/io/github/llmagentbuilder/core/AgentConfig.kt b/core/src/main/kotlin/io/github/llmagentbuilder/core/AgentConfig.kt
@@ -47,9 +47,20 @@ class TracingConfig {
     var exporter: TracingExporterConfig? = null
 }
 
+class MetricsExporterConfig {
+    lateinit var endpoint: String
+    var headers: Map<String, String>? = null
+}
+
+class MetricsConfig {
+    var enabled: Boolean? = false
+    var exporter: MetricsExporterConfig? = null
+}
+
 class ObservationConfig {
     var enabled: Boolean? = false
     var tracing: TracingConfig? = null
+    var metrics: MetricsConfig? = null
 }
 
 class AgentConfig {

diff --git a/core/src/main/kotlin/io/github/llmagentbuilder/core/ObservationPlugin.kt b/core/src/main/kotlin/io/github/llmagentbuilder/core/ObservationPlugin.kt
@@ -5,6 +5,6 @@ import io.micrometer.observation.ObservationRegistry
 interface ObservationPlugin {
     fun install(
         agentConfig: AgentConfig,
-        observationRegistry: ObservationRegistry
+        observationRegistry: ObservationRegistry,
     )
 }