KG-217. Implement span adapters for Langfuse and Weave clients

sdubov · sdubov · commit 0c92e85a8ebf · 2025-08-13T14:26:54.000+02:00
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/extension/SpanExt.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/extension/SpanExt.kt
@@ -2,7 +2,9 @@ package ai.koog.agents.features.opentelemetry.extension
 
 import ai.koog.agents.features.opentelemetry.attribute.Attribute
 import ai.koog.agents.features.opentelemetry.attribute.toSdkAttributes
+import ai.koog.agents.features.opentelemetry.event.EventBodyField
 import ai.koog.agents.features.opentelemetry.event.GenAIAgentEvent
+import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
 import ai.koog.agents.features.opentelemetry.span.SpanEndStatus
 import io.opentelemetry.api.trace.Span
 import io.opentelemetry.api.trace.SpanBuilder
@@ -36,3 +38,15 @@ internal fun Span.setEvents(events: List<GenAIAgentEvent>) {
         addEvent(event.name, attributes.toSdkAttributes())
     }
 }
+
+internal inline fun <reified TBodyField> GenAIAgentSpan.eventBodyFieldToAttribute(
+    event: GenAIAgentEvent,
+    attributeCreate: (TBodyField) -> Attribute
+) where TBodyField : EventBodyField {
+    event.bodyFields.filterIsInstance<TBodyField>().forEach { bodyField ->
+        val attributeFromEvent = attributeCreate(bodyField)
+        this.addAttribute(attributeFromEvent)
+    }
+
+    this.removeEvent(event)
+}
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetryConfig.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetryConfig.kt
@@ -2,7 +2,7 @@ package ai.koog.agents.features.opentelemetry.feature
 
 import ai.koog.agents.core.feature.config.FeatureConfig
 import ai.koog.agents.features.opentelemetry.attribute.addAttributes
-import ai.koog.agents.features.opentelemetry.integrations.SpanAdapter
+import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/SpanAdapter.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/SpanAdapter.kt
@@ -1,4 +1,4 @@
-package ai.koog.agents.features.opentelemetry.integrations
+package ai.koog.agents.features.opentelemetry.integration
 
 import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
 
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/langfuse/Langfuse.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/langfuse/Langfuse.kt
@@ -1,4 +1,4 @@
-package ai.koog.agents.features.opentelemetry.integrations
+package ai.koog.agents.features.opentelemetry.integration.langfuse
 
 import ai.koog.agents.features.opentelemetry.feature.OpenTelemetryConfig
 import io.github.oshai.kotlinlogging.KotlinLogging
@@ -11,7 +11,7 @@ import kotlin.time.Duration.Companion.seconds
 /**
  * Configure an OpenTelemetry span exporter that sends data to [Langfuse](https://langfuse.com/).
  *
- * @param langfuseUrl the base URL of the Langfuse instance. If not set is retrieved from `LANGFUSE_URL` environment variable. Defaults to [https://cloud.langfuse.com](https://cloud.langfuse.com).
+ * @param langfuseUrl the base URL of the Langfuse instance. If not set is retrieved from `LANGFUSE_HOST` environment variable. Defaults to [https://cloud.langfuse.com](https://cloud.langfuse.com).
  * @param langfusePublicKey if not set is retrieved from `LANGFUSE_PUBLIC_KEY` environment variable.
  * @param langfuseSecretKey if not set is retrieved from `LANGFUSE_SECRET_KEY` environment variable.
  * @param timeout OpenTelemetry SpanExporter timeout. See [io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporterBuilder.setTimeout].
@@ -46,7 +46,7 @@ public fun OpenTelemetryConfig.addLangfuseExporter(
             .build()
     )
 
-    // TODO: Add span adapter for Langfuse
+    addSpanAdapter(LangfuseSpanAdapter)
 }
 
 private val logger = KotlinLogging.logger {}
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/langfuse/LangfuseSpanAdapter.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/langfuse/LangfuseSpanAdapter.kt
@@ -0,0 +1,62 @@
+package ai.koog.agents.features.opentelemetry.integration.langfuse
+
+import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.features.opentelemetry.attribute.CustomAttribute
+import ai.koog.agents.features.opentelemetry.event.AssistantMessageEvent
+import ai.koog.agents.features.opentelemetry.event.ChoiceEvent
+import ai.koog.agents.features.opentelemetry.event.EventBodyFields
+import ai.koog.agents.features.opentelemetry.event.UserMessageEvent
+import ai.koog.agents.features.opentelemetry.extension.eventBodyFieldToAttribute
+import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
+import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
+import ai.koog.agents.features.opentelemetry.span.InferenceSpan
+
+@OptIn(InternalAgentsApi::class)
+internal object LangfuseSpanAdapter : SpanAdapter() {
+
+    override fun onBeforeSpanFinished(span: GenAIAgentSpan) {
+        when (span) {
+            is InferenceSpan -> span.prepareSpanAttributes()
+        }
+    }
+
+    //region Private Methods
+
+    private fun InferenceSpan.prepareSpanAttributes() {
+        this.events.forEach { event ->
+            when (event) {
+                is AssistantMessageEvent -> {
+                    this.eventBodyFieldToAttribute<EventBodyFields.Role>(event) { role ->
+                        CustomAttribute("gen_ai.completion.${role.key}", role.value)
+                    }
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.content", content.value)
+                    }
+                }
+
+                is ChoiceEvent -> {
+                    val index = event.index
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Role>(event) { role ->
+                        CustomAttribute("gen_ai.$index.${role.key}", role.value)
+                    }
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.$index.content", content.value)
+                    }
+                }
+
+                is UserMessageEvent -> {
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.content", content.value)
+                    }
+                }
+
+                else -> { }
+            }
+        }
+    }
+
+    //endregion Private Methods
+}
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/weave/Weave.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/weave/Weave.kt
@@ -1,4 +1,4 @@
-package ai.koog.agents.features.opentelemetry.integrations
+package ai.koog.agents.features.opentelemetry.integration.weave
 
 import ai.koog.agents.features.opentelemetry.feature.OpenTelemetryConfig
 import io.github.oshai.kotlinlogging.KotlinLogging
@@ -46,7 +46,7 @@ public fun OpenTelemetryConfig.addWeaveExporter(
             .build()
     )
 
-    // TODO: Add span adapter for Langfuse
+    addSpanAdapter(WeaveSpanAdapter)
 }
 
 private val logger = KotlinLogging.logger { }
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/weave/WeaveSpanAdapter.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmMain/kotlin/ai/koog/agents/features/opentelemetry/integration/weave/WeaveSpanAdapter.kt
@@ -0,0 +1,62 @@
+package ai.koog.agents.features.opentelemetry.integration.weave
+
+import ai.koog.agents.core.annotation.InternalAgentsApi
+import ai.koog.agents.features.opentelemetry.attribute.CustomAttribute
+import ai.koog.agents.features.opentelemetry.event.AssistantMessageEvent
+import ai.koog.agents.features.opentelemetry.event.ChoiceEvent
+import ai.koog.agents.features.opentelemetry.event.EventBodyFields
+import ai.koog.agents.features.opentelemetry.event.UserMessageEvent
+import ai.koog.agents.features.opentelemetry.extension.eventBodyFieldToAttribute
+import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
+import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
+import ai.koog.agents.features.opentelemetry.span.InferenceSpan
+
+@OptIn(InternalAgentsApi::class)
+internal object WeaveSpanAdapter : SpanAdapter() {
+
+    override fun onBeforeSpanFinished(span: GenAIAgentSpan) {
+        when (span) {
+            is InferenceSpan -> span.prepareSpanAttributes()
+        }
+    }
+
+    //region Private Methods
+
+    private fun InferenceSpan.prepareSpanAttributes() {
+        this.events.forEach { event ->
+            when (event) {
+                is AssistantMessageEvent -> {
+                    this.eventBodyFieldToAttribute<EventBodyFields.Role>(event) { role ->
+                        CustomAttribute("gen_ai.completion.${role.key}", role.value)
+                    }
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.content", content.value)
+                    }
+                }
+
+                is ChoiceEvent -> {
+                    val index = event.index
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Role>(event) { role ->
+                        CustomAttribute("gen_ai.$index.${role.key}", role.value)
+                    }
+
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.$index.content", content.value)
+                    }
+                }
+
+                is UserMessageEvent -> {
+                    this.eventBodyFieldToAttribute<EventBodyFields.Content>(event) { content ->
+                        CustomAttribute("gen_ai.completion.content", content.value)
+                    }
+                }
+
+                else -> { }
+            }
+        }
+    }
+
+    //endregion Private Methods
+}
diff --git a/agents/agents-features/agents-features-opentelemetry/src/jvmTest/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetryTest.kt b/agents/agents-features/agents-features-opentelemetry/src/jvmTest/kotlin/ai/koog/agents/features/opentelemetry/feature/OpenTelemetryTest.kt
@@ -12,7 +12,7 @@ import ai.koog.agents.features.opentelemetry.OpenTelemetryTestAPI.createAgent
 import ai.koog.agents.features.opentelemetry.attribute.CustomAttribute
 import ai.koog.agents.features.opentelemetry.attribute.SpanAttributes
 import ai.koog.agents.features.opentelemetry.attribute.SpanAttributes.Response.FinishReasonType
-import ai.koog.agents.features.opentelemetry.integrations.SpanAdapter
+import ai.koog.agents.features.opentelemetry.integration.SpanAdapter
 import ai.koog.agents.features.opentelemetry.mock.MockSpanExporter
 import ai.koog.agents.features.opentelemetry.mock.TestGetWeatherTool
 import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan
diff --git a/examples/src/main/kotlin/ai/koog/agents/example/features/langfuse/Langfuse.kt b/examples/src/main/kotlin/ai/koog/agents/example/features/langfuse/Langfuse.kt
@@ -3,7 +3,7 @@ package ai.koog.agents.example.features.langfuse
 import ai.koog.agents.core.agent.AIAgent
 import ai.koog.agents.example.ApiKeyService
 import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry
-import ai.koog.agents.features.opentelemetry.integrations.addLangfuseExporter
+import ai.koog.agents.features.opentelemetry.integration.langfuse.addLangfuseExporter
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
 import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter
diff --git a/examples/src/main/kotlin/ai/koog/agents/example/features/weave/Weave.kt b/examples/src/main/kotlin/ai/koog/agents/example/features/weave/Weave.kt
@@ -3,7 +3,7 @@ package ai.koog.agents.example.features.weave
 import ai.koog.agents.core.agent.AIAgent
 import ai.koog.agents.example.ApiKeyService
 import ai.koog.agents.features.opentelemetry.feature.OpenTelemetry
-import ai.koog.agents.features.opentelemetry.integrations.addWeaveExporter
+import ai.koog.agents.features.opentelemetry.integration.weave.addWeaveExporter
 import ai.koog.prompt.executor.clients.openai.OpenAIModels
 import ai.koog.prompt.executor.llms.all.simpleOpenAIExecutor
 import io.opentelemetry.exporter.otlp.http.trace.OtlpHttpSpanExporter

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package ai.koog.agents.features.opentelemetry.integrations`
	`1`	`+package ai.koog.agents.features.opentelemetry.integration`
`2`	`2`
`3`	`3`	`import ai.koog.agents.features.opentelemetry.span.GenAIAgentSpan`
`4`	`4`