JBAI-19233 Split tracing of the LLM clients' OkHttp handlers

examples/build.gradle.kts

@@ -18,6 +18,7 @@ dependencies {
     implementation(project(":tracing:core"))
     implementation(project(":tracing:gemini"))
     implementation(project(":tracing:ktor"))
+    implementation(project(":tracing:okhttp"))
     implementation(project(":tracing:openai"))
     testImplementation(libs.kotlin.test)
 }

examples/src/main/kotlin/ai/jetbrains/tracy/examples/clients/OkHttpClientAutotracingExample.kt

@@ -1,13 +1,13 @@
 package ai.jetbrains.tracy.examples.clients
 
-import ai.jetbrains.tracy.core.OpenTelemetryOkHttpInterceptor
 import ai.jetbrains.tracy.core.exporters.ConsoleExporterConfig
-import ai.jetbrains.tracy.core.instrument
 import ai.jetbrains.tracy.core.tracing.TracingManager
 import ai.jetbrains.tracy.core.tracing.configureOpenTelemetrySdk
 import ai.jetbrains.tracy.openai.adapters.OpenAILLMTracingAdapter
 import ai.jetbrains.tracy.gemini.adapters.GeminiLLMTracingAdapter
 import ai.jetbrains.tracy.anthropic.adapters.AnthropicLLMTracingAdapter
+import ai.jetbrains.tracy.okhttp.interceptors.OpenTelemetryOkHttpInterceptor
+import ai.jetbrains.tracy.okhttp.instrument
 import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.JsonPrimitive
 import kotlinx.serialization.json.buildJsonArray
@@ -40,6 +40,7 @@ fun main() {
     TracingManager.traceSensitiveContent()
 
     val apiToken = System.getenv("OPENAI_API_KEY") ?: error("Environment variable 'OPENAI_API_KEY' is not set")
+
     val requestBodyJson = buildJsonObject {
         put("model", JsonPrimitive("gpt-4o-mini"))
         put("messages", buildJsonArray {
@@ -50,18 +51,23 @@ fun main() {
         })
         put("temperature", JsonPrimitive(1.0))
     }
+
     val client = OkHttpClient()
     val instrumentedClient = instrument(client, OpenAILLMTracingAdapter())
+
     val requestBody = Json { prettyPrint = true }
         .encodeToString(requestBodyJson)
         .toRequestBody("application/json".toMediaType())
+
     val request = Request.Builder().url("https://api.openai.com/v1/chat/completions")
         .addHeader("Authorization", "Bearer $apiToken")
         .addHeader("Content-Type", "application/json")
         .post(requestBody)
         .build()
+
     instrumentedClient.newCall(request).execute().use { response ->
         println("Result: ${response.body?.string() ?: "<empty response>"}\nSee trace details in the console.")
     }
+
     TracingManager.flushTraces()
 }

settings.gradle.kts

@@ -22,6 +22,7 @@ include("tracing:anthropic")
 include("tracing:core")
 include("tracing:gemini")
 include("tracing:ktor")
+include("tracing:okhttp")
 include("tracing:openai")
 include("tracing:test-utils")
-includeBuild("plugin")
+includeBuild("plugin")

tracing/anthropic/build.gradle.kts

@@ -25,6 +25,7 @@ kotlin {
 
         jvmMain {
             dependencies {
+                implementation(project(":tracing:okhttp"))
                 implementation(libs.anthropic)
                 implementation(libs.okhttp)
                 implementation(libs.opentelemetry)

tracing/anthropic/src/jvmMain/kotlin/ai/jetbrains/tracy/anthropic/clients/AnthropicAIClient.kt

@@ -1,8 +1,8 @@
 package ai.jetbrains.tracy.anthropic.clients
 
-import ai.jetbrains.tracy.core.OpenTelemetryOkHttpInterceptor
-import ai.jetbrains.tracy.core.patchOpenAICompatibleClient
 import ai.jetbrains.tracy.anthropic.adapters.AnthropicLLMTracingAdapter
+import ai.jetbrains.tracy.okhttp.interceptors.OpenTelemetryOkHttpInterceptor
+import ai.jetbrains.tracy.okhttp.interceptors.patchOpenAICompatibleClient
 import com.anthropic.client.AnthropicClient
 
 fun instrument(client: AnthropicClient): AnthropicClient {

tracing/anthropic/src/jvmTest/kotlin/ai/jetbrains/tracy/anthropic/BaseAnthropicTracingTest.kt

@@ -1,7 +1,7 @@
 package ai.jetbrains.tracy.anthropic
 
-import ai.jetbrains.tracy.core.getFieldValue
-import ai.jetbrains.tracy.core.setFieldValue
+import ai.jetbrains.tracy.okhttp.interceptors.getFieldValue
+import ai.jetbrains.tracy.okhttp.interceptors.setFieldValue
 import ai.jetbrains.tracy.test.utils.BaseAITracingTest
 import com.anthropic.client.AnthropicClient
 import com.anthropic.client.okhttp.AnthropicOkHttpClient

tracing/core/src/jvmMain/kotlin/ai/jetbrains/tracy/core/http/protocol/Request.kt

@@ -43,8 +43,6 @@ sealed class RequestBody {
     object Empty : RequestBody()
 }
 
-fun MediaType.toContentType(): ContentType = ContentType.parse(this.toString())
-
 fun RequestBody.asJson(): JsonElement? {
     return when (this) {
         is RequestBody.Json -> this.json
@@ -100,4 +98,6 @@ fun ByteArray.asRequestBody(contentType: ContentType): RequestBody? {
         }
         else -> null
     }
-}
+}
+
+fun MediaType.toContentType() = ContentType.parse(this.toString())

tracing/core/src/jvmMain/kotlin/ai/jetbrains/tracy/core/http/protocol/Url.kt

@@ -1,17 +1,7 @@
 package ai.jetbrains.tracy.core.http.protocol
 
-import io.ktor.http.URLBuilder
-import io.ktor.http.Url as KtorUrl
-import okhttp3.HttpUrl
-
 data class Url(
     val scheme: String,
     val host: String,
     val pathSegments: List<String>,
 )
-
-fun HttpUrl.toProtocolUrl() = Url(scheme, host, pathSegments)
-
-fun URLBuilder.toProtocolUrl() = Url(protocol.name, host, pathSegments)
-
-fun KtorUrl.toProtocolUrl() = Url(protocol.name, host, segments)

tracing/gemini/build.gradle.kts

@@ -18,6 +18,7 @@ kotlin {
         commonMain {
             dependencies {
                 api(project(":tracing:core"))
+                api(project(":tracing:okhttp"))
                 implementation(libs.kotlinx.serialization.core)
                 implementation(libs.kotlinx.serialization.json)
             }

tracing/gemini/src/jvmMain/kotlin/ai/jetbrains/tracy/gemini/clients/GeminiClient.kt

@@ -1,8 +1,8 @@
 package ai.jetbrains.tracy.gemini.clients
 
-import ai.jetbrains.tracy.core.OpenTelemetryOkHttpInterceptor
 import ai.jetbrains.tracy.gemini.adapters.GeminiLLMTracingAdapter
-import ai.jetbrains.tracy.core.patchInterceptors
+import ai.jetbrains.tracy.okhttp.interceptors.OpenTelemetryOkHttpInterceptor
+import ai.jetbrains.tracy.okhttp.interceptors.patchInterceptors
 import okhttp3.Interceptor
 import okhttp3.OkHttpClient
 import com.google.genai.Client as GeminiClient

tracing/ktor/src/jvmMain/kotlin/ai/jetbrains/tracy/ktor/KtorHttpClient.kt

@@ -1,9 +1,14 @@
 package ai.jetbrains.tracy.ktor
 
 import ai.jetbrains.tracy.core.adapters.LLMTracingAdapter
+import ai.jetbrains.tracy.core.http.protocol.Request
+import ai.jetbrains.tracy.core.http.protocol.RequestBody
+import ai.jetbrains.tracy.core.http.protocol.Response
+import ai.jetbrains.tracy.core.http.protocol.ResponseBody
+import ai.jetbrains.tracy.core.tracing.TracingManager
 import ai.jetbrains.tracy.core.fluent.processor.Span
 import ai.jetbrains.tracy.core.http.protocol.*
-import ai.jetbrains.tracy.core.tracing.TracingManager
+import ai.jetbrains.tracy.ktor.extensions.toProtocolUrl
 import io.ktor.client.*
 import io.ktor.client.plugins.api.*
 import io.ktor.client.request.*

tracing/ktor/src/jvmMain/kotlin/ai/jetbrains/tracy/ktor/extensions/Extensions.kt

@@ -0,0 +1,10 @@
+package ai.jetbrains.tracy.ktor.extensions
+
+import ai.jetbrains.tracy.core.http.protocol.Url
+import io.ktor.http.URLBuilder
+import io.ktor.http.Url as KtorUrl
+
+
+fun URLBuilder.toProtocolUrl() = Url(protocol.name, host, pathSegments)
+
+fun KtorUrl.toProtocolUrl() = Url(protocol.name, host, segments)

tracing/okhttp/build.gradle.kts

@@ -0,0 +1,48 @@
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget.JVM_17
+
+plugins {
+    kotlin("multiplatform")
+    alias(libs.plugins.kotlin.serialization)
+    id("ai.jetbrains.tracy.published-artifact")
+    id("ai.kotlin.dokka")
+}
+
+kotlin {
+    jvmToolchain(17)
+
+    jvm {
+        compilerOptions.jvmTarget = JVM_17
+    }
+
+    sourceSets {
+        commonMain {
+            dependencies {
+                api(project(":tracing:core"))
+                implementation(libs.kotlinx.serialization.core)
+                implementation(libs.kotlinx.serialization.json)
+            }
+        }
+
+        jvmMain {
+            dependencies {
+                implementation(libs.kotlin.reflect)
+                implementation(libs.okhttp)
+                implementation(libs.ktor.client)
+                implementation(libs.opentelemetry)
+                implementation(libs.opentelemetry.kotlin)
+                implementation(libs.opentelemetry.sdk)
+                implementation(libs.opentelemetry.semconv.incubating)
+                implementation(libs.kotlin.logging)
+            }
+        }
+
+        jvmTest {
+            dependencies {
+                implementation(libs.kotlin.test)
+                implementation(libs.junit.params)
+                implementation(libs.kotlinx.coroutines.test)
+                implementation(libs.opentelemetry.sdk.testing)
+            }
+        }
+    }
+}

tracing/okhttp/src/jvmMain/kotlin/ai/jetbrains/tracy/okhttp/OkHttpClient.kt

@@ -0,0 +1,15 @@
+package ai.jetbrains.tracy.okhttp
+
+import ai.jetbrains.tracy.core.adapters.LLMTracingAdapter
+import ai.jetbrains.tracy.okhttp.interceptors.OpenTelemetryOkHttpInterceptor
+import ai.jetbrains.tracy.okhttp.interceptors.patchInterceptorsInplace
+import okhttp3.OkHttpClient
+
+fun instrument(client: OkHttpClient, adapter: LLMTracingAdapter): OkHttpClient {
+    val clientBuilder = client.newBuilder()
+
+    val interceptor = OpenTelemetryOkHttpInterceptor(adapter)
+    patchInterceptorsInplace(clientBuilder.interceptors(), interceptor)
+
+    return clientBuilder.build()
+}

tracing/okhttp/src/jvmMain/kotlin/ai/jetbrains/tracy/okhttp/extensions/Extensions.kt

@@ -0,0 +1,6 @@
+package ai.jetbrains.tracy.okhttp.extensions
+
+import ai.jetbrains.tracy.core.http.protocol.Url
+import okhttp3.HttpUrl
+
+fun HttpUrl.toProtocolUrl() = Url(scheme, host, pathSegments)

tracing/okhttp/src/jvmMain/kotlin/ai/jetbrains/tracy/okhttp/interceptors/FieldsPatching.kt

@@ -0,0 +1,63 @@
+package ai.jetbrains.tracy.okhttp.interceptors
+
+/**
+ * Retrieves the value of a specified field from the given instance (usually, **HTTP clients**),
+ * including fields declared in superclasses.
+ *
+ * @param instance The object instance from which to retrieve the field value (**an HTTP client**).
+ * @param fieldName The name of the field to retrieve.
+ * @return The value of the specified field.
+ * @throws NoSuchFieldException If the field with the specified name is not found.
+ * @throws IllegalStateException If the specified field is found but its value is null.
+ *
+ * @see setFieldValue
+ */
+fun getFieldValue(instance: Any, fieldName: String): Any {
+    var cls: Class<*>? = instance.javaClass
+    while (cls != null) {
+        try {
+            val field = cls.getDeclaredField(fieldName)
+            field.isAccessible = true
+            return field.get(instance) ?: throw IllegalStateException("Field '$fieldName' is null")
+        } catch (_: NoSuchFieldException) {
+            cls = cls.superclass
+        }
+    }
+    throw NoSuchFieldException("Field '$fieldName' not found in ${instance.javaClass.name}")
+}
+
+/**
+ * Sets the value of a specified field on a given object instance (usually, **HTTP clients**).
+ *
+ * This method traverses the class hierarchy of the provided instance to find the
+ * specified field and sets its value. If the field is not found in the class or
+ * its superclasses, or if the value is `null` while the field type is primitive,
+ * an exception will be thrown.
+ *
+ * @param instance The object instance whose field value is to be modified (**an HTTP client**).
+ * @param fieldName The name of the field to be updated.
+ * @param value The new value to assign to the field. Can be `null` if the field type allows it.
+ * @throws IllegalArgumentException If an attempt is made to set a null value on a primitive field.
+ * @throws NoSuchFieldException If the specified field is not found in the class hierarchy of the instance.
+ *
+ * @see getFieldValue
+ */
+fun setFieldValue(instance: Any, fieldName: String, value: Any?) {
+    var cls: Class<*>? = instance.javaClass
+    while (cls != null) {
+        try {
+            val field = cls.getDeclaredField(fieldName)
+            field.isAccessible = true
+
+            if (value == null && field.type.isPrimitive) {
+                throw IllegalArgumentException("Cannot set primitive field '$fieldName' to null")
+            }
+
+            field.set(instance, value)
+            return
+        } catch (_: NoSuchFieldException) {
+            cls = cls.superclass
+        }
+    }
+    throw NoSuchFieldException("Field '$fieldName' not found in ${instance.javaClass.name}")
+}