diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/OpenTelemetryConfig.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/OpenTelemetryConfig.kt index 745e8c77..496090bc 100644 --- a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/OpenTelemetryConfig.kt +++ b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/OpenTelemetryConfig.kt @@ -6,6 +6,7 @@ import io.opentelemetry.sdk.OpenTelemetrySdk import io.opentelemetry.sdk.common.CompletableResultCode import io.opentelemetry.sdk.resources.Resource import io.opentelemetry.sdk.trace.SdkTracerProvider +import io.opentelemetry.sdk.trace.SpanProcessor import io.opentelemetry.sdk.trace.data.SpanData import io.opentelemetry.sdk.trace.export.BatchSpanProcessor import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor @@ -84,11 +85,13 @@ object OpenTelemetryConfig { } private fun buildOtlpEndpoint(baseUrl: String?): String { - var effectiveBaseUrl = baseUrl - if (effectiveBaseUrl.isNullOrEmpty()) effectiveBaseUrl = System.getenv("LANGSMITH_ENDPOINT") - if (effectiveBaseUrl.isNullOrEmpty()) effectiveBaseUrl = DEFAULT_BASE_URL - if (effectiveBaseUrl!!.endsWith("/")) effectiveBaseUrl = effectiveBaseUrl.dropLast(1) - return effectiveBaseUrl + OTLP_TRACES_PATH + val base = + (baseUrl?.takeIf { it.isNotBlank() } + ?: System.getenv("LANGSMITH_ENDPOINT")?.takeIf { it.isNotBlank() } + ?: DEFAULT_BASE_URL) + .trim() + .removeSuffix("/") + return base + OTLP_TRACES_PATH } class Builder { @@ -113,6 +116,31 @@ object OpenTelemetryConfig { fun maxBatchSize(maxBatchSize: Int) = apply { this.maxBatchSize = maxBatchSize } + fun buildSpanProcessor(): SpanProcessor { + require(!apiKey.isNullOrEmpty()) { + "LangSmith API key is required. Set it using apiKey() or LANGSMITH_API_KEY environment variable." + } + val endpointUrl = buildOtlpEndpoint(baseUrl) + val exporterBuilder = + OtlpHttpSpanExporter.builder() + .setEndpoint(endpointUrl) + .addHeader("x-api-key", apiKey!!) + if (!projectName.isNullOrEmpty()) { + exporterBuilder.addHeader("Langsmith-Project", projectName!!) + } + val spanExporter = exporterBuilder.build() + val loggingExporter = LoggingSpanExporter(spanExporter) + return when (processorType) { + SpanProcessorType.SIMPLE -> SimpleSpanProcessor.create(loggingExporter) + SpanProcessorType.BATCH -> + BatchSpanProcessor.builder(loggingExporter) + .setScheduleDelay(100, TimeUnit.MILLISECONDS) + .setMaxExportBatchSize(maxBatchSize) + .setExporterTimeout(5, TimeUnit.SECONDS) + .build() + } + } + fun build(): io.opentelemetry.api.OpenTelemetry { require(!apiKey.isNullOrEmpty()) { "LangSmith API key is required. Set it using apiKey() or LANGSMITH_API_KEY environment variable." @@ -169,39 +197,35 @@ object OpenTelemetryConfig { } } val result = delegate.export(spans) - if (DEBUG) { - try { - result.join(5, TimeUnit.SECONDS) - if (!result.isSuccess) { + // Always wait for export to complete so flush/shutdown don't run before the HTTP + // request finishes. Without this, shutdown can abort in-flight exports and traces are + // lost. + try { + result.join(5, TimeUnit.SECONDS) + } catch (e: Exception) { + if (DEBUG) logger.error("[LangSmith ERROR] Exception waiting for export result", e) + else logger.error("[LangSmith] Exception waiting for export result", e) + } + if (!result.isSuccess) { + logger.error( + "[LangSmith ERROR] Failed to export ${spans.size} span(s) to LangSmith" + ) + logExportException(result) + if (DEBUG) { + for (span in spans) { logger.error( - "[LangSmith ERROR] Failed to export ${spans.size} span(s) to LangSmith" + " - ${span.name} (traceId=${span.traceId}, spanId=${span.spanId})" ) - logExportException(result) - for (span in spans) { - logger.error( - " - ${span.name} (traceId=${span.traceId}, spanId=${span.spanId})" - ) - } - logger.error( - " This usually indicates a network error, authentication problem, or invalid span data" - ) - logger.error(" Check your LANGSMITH_API_KEY and network connectivity") - } else { - logger.debug("[LangSmith] Successfully exported ${spans.size} span(s)") - } - } catch (e: Exception) { - logger.error("[LangSmith ERROR] Exception waiting for export result", e) - } - } else { - result.whenComplete { - if (!result.isSuccess) { - logger.error( - "[LangSmith ERROR] Failed to export ${spans.size} span(s) to LangSmith" - ) - logExportException(result) - logger.error(" Set LANGSMITH_DEBUG=true for more details") } + logger.error( + " This usually indicates a network error, authentication problem, or invalid span data" + ) + logger.error(" Check your LANGSMITH_API_KEY and network connectivity") + } else { + logger.error(" Set LANGSMITH_DEBUG=true for more details") } + } else if (DEBUG) { + logger.debug("[LangSmith] Successfully exported ${spans.size} span(s)") } return result } @@ -222,6 +246,7 @@ object OpenTelemetryConfig { } companion object { + private val logger = LoggerFactory.getLogger(LoggingSpanExporter::class.java) private val DEBUG = java.lang.Boolean.getBoolean("langsmith.debug") || "true".equals(System.getenv("LANGSMITH_DEBUG"), ignoreCase = true) diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/TracingUtils.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/TracingUtils.kt index 4b41bd70..481fa427 100644 --- a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/TracingUtils.kt +++ b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/TracingUtils.kt @@ -1,5 +1,7 @@ package com.langchain.smith.wrappers.openai +import com.fasterxml.jackson.databind.ObjectMapper +import io.opentelemetry.api.common.AttributeKey import io.opentelemetry.api.trace.Span import io.opentelemetry.api.trace.SpanBuilder import io.opentelemetry.api.trace.SpanKind @@ -8,26 +10,10 @@ import io.opentelemetry.api.trace.Tracer /** Internal utility for OpenTelemetry span creation and management. */ internal object TracingUtils { private const val INSTRUMENTATION_NAME = "langsmith-java-otel-wrappers" + private val jsonMapper = ObjectMapper() - fun getTracer(): Tracer { - return try { - val tracer = - io.opentelemetry.api.GlobalOpenTelemetry.get().getTracer(INSTRUMENTATION_NAME) - val debug = - java.lang.Boolean.getBoolean("langsmith.debug") || - "true".equals(System.getenv("LANGSMITH_DEBUG"), ignoreCase = true) - if (debug) { - val otel = io.opentelemetry.api.GlobalOpenTelemetry.get() - val isNoop = otel.javaClass.name.contains("Noop") - println( - "[TracingUtils] Tracer obtained: ${tracer.javaClass.name}, OpenTelemetry isNoop: $isNoop" - ) - } - tracer - } catch (e: Exception) { - io.opentelemetry.api.GlobalOpenTelemetry.get().getTracer(INSTRUMENTATION_NAME) - } - } + fun getTracer(): Tracer = + io.opentelemetry.api.GlobalOpenTelemetry.get().getTracer(INSTRUMENTATION_NAME) fun createSpanBuilder( model: String?, @@ -40,29 +26,33 @@ internal object TracingUtils { tracer .spanBuilder(spanName) .setSpanKind(SpanKind.CLIENT) - .setAttribute("gen_ai.system", "openai") - .setAttribute("gen_ai.operation.name", operationType) - .setAttribute("gen_ai.provider.name", "openai") - spanKind?.let { builder.setAttribute("langsmith.span.kind", it) } + .setAttribute(AttributeKey.stringKey("gen_ai.system"), "openai") + .setAttribute(AttributeKey.stringKey("gen_ai.operation.name"), operationType) + .setAttribute(AttributeKey.stringKey("gen_ai.provider.name"), "openai") + spanKind?.let { builder.setAttribute(AttributeKey.stringKey("langsmith.span.kind"), it) } return builder } fun setRequestAttributes(span: Span, model: String?) { - model?.let { span.setAttribute("gen_ai.request.model", it) } + model?.let { span.setAttribute(AttributeKey.stringKey("gen_ai.request.model"), it) } } fun setRequestParameters(span: Span, temperature: Double?, topP: Double?, maxTokens: Long?) { - temperature?.let { span.setAttribute("gen_ai.request.temperature", it) } - topP?.let { span.setAttribute("gen_ai.request.top_p", it) } - maxTokens?.let { span.setAttribute("gen_ai.request.max_tokens", it) } + temperature?.let { + span.setAttribute(AttributeKey.doubleKey("gen_ai.request.temperature"), it) + } + topP?.let { span.setAttribute(AttributeKey.doubleKey("gen_ai.request.top_p"), it) } + maxTokens?.let { span.setAttribute(AttributeKey.longKey("gen_ai.request.max_tokens"), it) } } fun setInputMessages(span: Span, messagesJson: String?) { - messagesJson?.let { span.setAttribute("gen_ai.input.messages", it) } + messagesJson?.let { span.setAttribute(AttributeKey.stringKey("gen_ai.input.messages"), it) } } fun setOutputMessages(span: Span, messagesJson: String?) { - messagesJson?.let { span.setAttribute("gen_ai.output.messages", it) } + messagesJson?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.output.messages"), it) + } } fun setResponseAttributes( @@ -71,27 +61,30 @@ internal object TracingUtils { outputTokens: Long?, totalTokens: Long?, ) { - inputTokens?.let { span.setAttribute("gen_ai.usage.input_tokens", it) } - outputTokens?.let { span.setAttribute("gen_ai.usage.output_tokens", it) } - totalTokens?.let { span.setAttribute("gen_ai.usage.total_tokens", it) } + inputTokens?.let { + span.setAttribute(AttributeKey.longKey("gen_ai.usage.input_tokens"), it) + } + outputTokens?.let { + span.setAttribute(AttributeKey.longKey("gen_ai.usage.output_tokens"), it) + } + totalTokens?.let { + span.setAttribute(AttributeKey.longKey("gen_ai.usage.total_tokens"), it) + } } fun setResponseMetadata(span: Span, responseModel: String?, finishReason: String?) { - responseModel?.let { span.setAttribute("gen_ai.response.model", it) } - finishReason?.let { span.setAttribute("gen_ai.response.finish_reason", it) } + responseModel?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.response.model"), it) + } + finishReason?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.response.finish_reason"), it) + } } fun recordException(span: Span, exception: Throwable) { span.recordException(exception) - span.setAttribute("error", true) + span.setAttribute(AttributeKey.booleanKey("error"), true) } - fun escapeJsonString(str: String?): String { - if (str == null) return "" - return str.replace("\\", "\\\\") - .replace("\"", "\\\"") - .replace("\n", "\\n") - .replace("\r", "\\r") - .replace("\t", "\\t") - } + fun writeJson(value: Any): String = jsonMapper.writeValueAsString(value) } diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedChatService.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedChatService.kt index 439ea6c7..42c9605b 100644 --- a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedChatService.kt +++ b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedChatService.kt @@ -3,9 +3,11 @@ package com.langchain.smith.wrappers.openai import com.openai.core.ClientOptions import com.openai.core.RequestOptions import com.openai.core.http.StreamResponse +import com.openai.helpers.ChatCompletionAccumulator import com.openai.models.chat.completions.ChatCompletion import com.openai.models.chat.completions.ChatCompletionChunk import com.openai.models.chat.completions.ChatCompletionCreateParams +import com.openai.models.chat.completions.ChatCompletionStreamOptions import com.openai.models.chat.completions.StructuredChatCompletion import com.openai.models.chat.completions.StructuredChatCompletionCreateParams import com.openai.services.blocking.ChatService @@ -14,7 +16,8 @@ import io.opentelemetry.api.common.AttributeKey import io.opentelemetry.api.trace.Span import java.util.function.Consumer import java.util.regex.Pattern -import org.slf4j.LoggerFactory +import java.util.stream.Stream +import kotlin.jvm.optionals.getOrNull /** Wrapped ChatService that adds OpenTelemetry tracing to chat completion operations. */ internal class WrappedChatService(private val delegate: ChatService) : ChatService { @@ -29,9 +32,6 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi private class WrappedChatCompletionService(private val delegate: ChatCompletionService) : ChatCompletionService { - companion object { - private val logger = LoggerFactory.getLogger(WrappedChatCompletionService::class.java) - } override fun withRawResponse() = delegate.withRawResponse() @@ -54,14 +54,6 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi ): ChatCompletion { val model = params.model()?.toString() val span = TracingUtils.createSpanBuilder(model, "chat").startSpan() - if (logger.isDebugEnabled) { - logger.debug( - "[WrappedChatService] Created span: {}, isRecording: {}, traceId: {}", - span.spanContext.spanId, - span.isRecording, - span.spanContext.traceId, - ) - } try { span.makeCurrent().use { setExperimentContextAttributes(span) @@ -72,41 +64,50 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi params.topP().orElse(null), params.maxCompletionTokens().orElse(null), ) - formatInputMessages(params).let { TracingUtils.setInputMessages(span, it) } - extractPromptFromParams(params) - ?.takeIf { it.isNotEmpty() } - ?.let { span.setAttribute(AttributeKey.stringKey("gen_ai.prompt"), it) } + val promptTextNonStream = + extractPromptFromParams(params)?.takeIf { it.isNotEmpty() } + promptTextNonStream?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.prompt"), it) + } + TracingUtils.setInputMessages(span, formatInputMessages(params)) val result = if (requestOptions == null) delegate.create(params) else delegate.create(params, requestOptions) - val responseModel = result.model() - val finishReason = - result.choices().firstOrNull()?.finishReason()?.toString() ?: "stop" - TracingUtils.setResponseMetadata(span, responseModel, finishReason) - formatOutputMessages(result).let { TracingUtils.setOutputMessages(span, it) } - extractCompletionFromResult(result) - ?.takeIf { it.isNotEmpty() } - ?.let { span.setAttribute(AttributeKey.stringKey("gen_ai.completion"), it) } - result.usage().ifPresent { usage -> - TracingUtils.setResponseAttributes( - span, - usage.promptTokens().toLong(), - usage.completionTokens().toLong(), - usage.totalTokens().toLong(), - ) - } + applyChatCompletionToSpan(span, result) return result } } catch (e: Exception) { TracingUtils.recordException(span, e) throw e } finally { - if (logger.isDebugEnabled) - logger.debug("[WrappedChatService] Ending span: {}", span.spanContext.spanId) span.end() } } + /** + * Applies completion output, gen_ai.completion JSON, and usage to the span (shared by + * non-streaming and streaming). + */ + private fun applyChatCompletionToSpan(span: Span, completion: ChatCompletion) { + formatOutputMessages(completion).let { TracingUtils.setOutputMessages(span, it) } + span.setAttribute( + AttributeKey.stringKey("gen_ai.completion"), + "{\"messages\":${formatOutputMessages(completion)}}", + ) + val responseModel = completion.model() + val finishReason = + completion.choices().firstOrNull()?.finishReason()?.toString() ?: "stop" + TracingUtils.setResponseMetadata(span, responseModel, finishReason) + completion.usage().ifPresent { usage -> + TracingUtils.setResponseAttributes( + span, + usage.promptTokens().toLong(), + usage.completionTokens().toLong(), + usage.totalTokens().toLong(), + ) + } + } + private fun setExperimentContextAttributes(span: Span) { ExperimentContext.current() .getReferenceExampleId() @@ -127,17 +128,8 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi private fun formatInputMessages(params: ChatCompletionCreateParams): String { if (params.messages().isEmpty()) return "[]" - if (logger.isDebugEnabled) - logger.debug( - "[formatInputMessages] Processing {} message(s)", - params.messages().size, - ) - val json = StringBuilder("[") - var first = true + val messages = mutableListOf>() for (messageParam in params.messages()) { - if (!first) json.append(",") - first = false - json.append("{") var role: String? = null var content: String? = null try { @@ -224,18 +216,18 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi if (content == null) { content = extractContentFromToString(messageParam.toString(), role) } - if (role != null) json.append("\"role\":\"").append(role).append("\"") - if (content != null) { - if (role != null) json.append(",") - json - .append("\"content\":\"") - .append(TracingUtils.escapeJsonString(content)) - .append("\"") - } - json.append("}") + val outRole = role ?: "user" + val partType = if (role == "tool") "tool_call_response" else "text" + val partKey = if (role == "tool") "response" else "content" + val parts = + if (content != null) { + listOf(mapOf("type" to partType, partKey to content)) + } else { + emptyList>() + } + messages.add(mapOf("role" to outRole, "parts" to parts)) } - json.append("]") - return json.toString() + return TracingUtils.writeJson(messages) } private fun extractContentFromToString(messageStr: String, role: String?): String? { @@ -282,79 +274,47 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi return null } + /** OTel GenAI schema: each message has role, parts array, and finish_reason (output). */ private fun formatOutputMessages(completion: ChatCompletion): String { if (completion.choices().isEmpty()) return "[]" - val json = StringBuilder("[") - var first = true - for (choice in completion.choices()) { - if (!first) json.append(",") - first = false - val message = choice.message() - json.append("{\"role\":\"assistant\"") - message.content().ifPresent { - json - .append(",\"content\":\"") - .append(TracingUtils.escapeJsonString(it)) - .append("\"") - } - message.toolCalls().ifPresent { toolCalls -> - if (toolCalls.isNotEmpty()) { - json.append(",\"tool_calls\":[") - var firstTc = true + val choices = + completion.choices().map { choice -> + val message = choice.message() + val finishReason = choice.finishReason()?.toString()?.lowercase() ?: "stop" + val parts = mutableListOf>() + message.content().ifPresent { content -> + parts.add(mapOf("type" to "text", "content" to content)) + } + message.toolCalls().ifPresent { toolCalls -> for (toolCall in toolCalls) { if (!toolCall.isFunction()) continue - if (!firstTc) json.append(",") - firstTc = false val fn = toolCall.asFunction() - json - .append("{\"id\":\"") - .append(TracingUtils.escapeJsonString(fn.id())) - .append("\"") - json.append(",\"type\":\"function\"") - json - .append(",\"function\":{\"name\":\"") - .append(TracingUtils.escapeJsonString(fn.function().name())) - json - .append("\",\"arguments\":\"") - .append(TracingUtils.escapeJsonString(fn.function().arguments())) - .append("\"}") - json.append("}") + parts.add( + mapOf( + "type" to "tool_call", + "id" to fn.id(), + "name" to fn.function().name(), + "arguments" to fn.function().arguments(), + ) + ) } - json.append("]") } + mapOf("role" to "assistant", "parts" to parts, "finish_reason" to finishReason) } - json.append("}") - } - json.append("]") - return json.toString() + return TracingUtils.writeJson(choices) } private fun extractPromptFromParams(params: ChatCompletionCreateParams): String? { for (messageParam in params.messages()) { - try { - if ( - messageParam.javaClass.getMethod("isUser").invoke(messageParam) as Boolean - ) { - val userMessage = - messageParam.javaClass.getMethod("asUser").invoke(messageParam) - val content = userMessage.javaClass.getMethod("content").invoke(userMessage) - when (content) { - is String -> return content - is List<*> -> - if (content.isNotEmpty()) - (content[0] as? String)?.let { - return it - } - } - } - } catch (_: Exception) {} + if (messageParam.isUser()) { + val userMessage = messageParam.asUser() + val text = userMessage.content().text().getOrNull() + if (!text.isNullOrEmpty()) return text + } } return null } - private fun extractCompletionFromResult(result: ChatCompletion): String? = - result.choices().firstOrNull()?.message()?.content()?.orElse(null) - override fun create( params: StructuredChatCompletionCreateParams ): StructuredChatCompletion = createStructured(params, null) @@ -381,7 +341,12 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi raw.topP().orElse(null), raw.maxCompletionTokens().orElse(null), ) - formatInputMessages(raw).let { TracingUtils.setInputMessages(span, it) } + val promptStructured = + extractPromptFromParams(raw)?.takeIf { it.isNotEmpty() } + promptStructured?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.prompt"), it) + } + TracingUtils.setInputMessages(span, formatInputMessages(raw)) } val result = if (requestOptions == null) delegate.create(params) @@ -414,6 +379,22 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi requestOptions: RequestOptions, ): StreamResponse = createStreamingChat(params, requestOptions) + /** + * Ensures the request asks for usage in the stream so we get token counts in the final + * chunk. Without stream_options.include_usage=true the API does not send usage and traces + * lack token metadata. + */ + private fun ensureStreamIncludeUsage( + params: ChatCompletionCreateParams + ): ChatCompletionCreateParams { + val existing = params.streamOptions().orElse(null) + if (existing != null && existing.includeUsage().orElse(false)) return params + val streamOptsBuilder = ChatCompletionStreamOptions.builder().includeUsage(true) + existing?.includeObfuscation()?.ifPresent { streamOptsBuilder.includeObfuscation(it) } + val newOpts = streamOptsBuilder.build() + return params.toBuilder().streamOptions(newOpts).build() + } + private fun createStreamingChat( params: ChatCompletionCreateParams, requestOptions: RequestOptions?, @@ -431,15 +412,30 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi params.topP().orElse(null), params.maxCompletionTokens().orElse(null), ) - formatInputMessages(params).let { TracingUtils.setInputMessages(span, it) } - return if (requestOptions == null) delegate.createStreaming(params) - else delegate.createStreaming(params, requestOptions!!) + val promptText = extractPromptFromParams(params)?.takeIf { it.isNotEmpty() } + promptText?.let { + span.setAttribute(AttributeKey.stringKey("gen_ai.prompt"), it) + } + TracingUtils.setInputMessages(span, formatInputMessages(params)) + val accumulator = ChatCompletionAccumulator.create() + val paramsWithUsage = ensureStreamIncludeUsage(params) + val startNanos = System.nanoTime() + val delegateStream = + if (requestOptions == null) delegate.createStreaming(paramsWithUsage) + else delegate.createStreaming(paramsWithUsage, requestOptions!!) + return TracedChatStreamResponse( + span, + accumulator, + delegateStream, + startNanos, + ) { s, c -> + c?.let { applyChatCompletionToSpan(s, it) } + } } } catch (e: Exception) { TracingUtils.recordException(span, e) - throw e - } finally { span.end() + throw e } } @@ -524,3 +520,54 @@ internal class WrappedChatService(private val delegate: ChatService) : ChatServi ) = delegate.delete(params, requestOptions) } } + +private class TracedChatStreamResponse( + private val span: Span, + private val accumulator: ChatCompletionAccumulator, + private val delegate: StreamResponse, + private val startNanos: Long, + private val onComplete: (Span, ChatCompletion?) -> Unit, +) : StreamResponse { + + @Volatile private var closed = false + @Volatile private var timeToFirstTokenSet = false + + override fun stream(): Stream = + delegate.stream().map { chunk -> + if (!timeToFirstTokenSet) { + val hasContent = + chunk.choices().isNotEmpty() && + chunk.choices()[0].delta().content().getOrNull()?.isNotEmpty() == true + if (hasContent) { + val ttftMs = (System.nanoTime() - startNanos) / 1_000_000 + span.setAttribute( + AttributeKey.longKey("gen_ai.usage.time_to_first_token_ms"), + ttftMs, + ) + timeToFirstTokenSet = true + } + } + accumulator.accumulate(chunk) + chunk + } + + override fun close() { + if (closed) return + closed = true + try { + try { + val completion = accumulator.chatCompletion() + onComplete(span, completion) + } catch (_: IllegalStateException) { + // Stream closed before a complete response + } + } finally { + span.end() + try { + delegate.close() + } catch (_: Exception) { + // Already ended span; ensure delegate is closed best-effort + } + } + } +} diff --git a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedResponseService.kt b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedResponseService.kt index cf2ee28f..a9ca43fc 100644 --- a/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedResponseService.kt +++ b/langsmith-java-core/src/main/kotlin/com/langchain/smith/wrappers/openai/WrappedResponseService.kt @@ -3,6 +3,7 @@ package com.langchain.smith.wrappers.openai import com.openai.core.ClientOptions import com.openai.core.RequestOptions import com.openai.core.http.StreamResponse +import com.openai.helpers.ResponseAccumulator import com.openai.models.responses.Response import com.openai.models.responses.ResponseCreateParams import com.openai.models.responses.ResponseStreamEvent @@ -12,6 +13,8 @@ import com.openai.services.blocking.ResponseService import io.opentelemetry.api.common.AttributeKey import io.opentelemetry.api.trace.Span import java.util.function.Consumer +import java.util.stream.Stream +import kotlin.jvm.optionals.getOrNull /** Wrapped ResponseService that adds OpenTelemetry tracing to response operations. */ internal class WrappedResponseService(private val delegate: ResponseService) : ResponseService { @@ -175,18 +178,20 @@ internal class WrappedResponseService(private val delegate: ResponseService) : R null, ) } - return when { - params == null && requestOptions == null -> delegate.createStreaming() - params == null -> delegate.createStreaming(requestOptions!!) - requestOptions == null -> delegate.createStreaming(params) - else -> delegate.createStreaming(params, requestOptions) - } + val accumulator = ResponseAccumulator.create() + val delegateStream = + when { + params == null && requestOptions == null -> delegate.createStreaming() + params == null -> delegate.createStreaming(requestOptions!!) + requestOptions == null -> delegate.createStreaming(params) + else -> delegate.createStreaming(params, requestOptions) + } + return TracedResponseStreamResponse(span, accumulator, delegateStream) } } catch (e: Exception) { TracingUtils.recordException(span, e) - throw e - } finally { span.end() + throw e } } @@ -222,16 +227,18 @@ internal class WrappedResponseService(private val delegate: ResponseService) : R null, ) } - return when { - requestOptions == null -> delegate.createStreaming(params!!) - else -> delegate.createStreaming(params!!, requestOptions) - } + val accumulator = ResponseAccumulator.create() + val delegateStream = + when { + requestOptions == null -> delegate.createStreaming(params!!) + else -> delegate.createStreaming(params!!, requestOptions) + } + return TracedResponseStreamResponse(span, accumulator, delegateStream) } } catch (e: Exception) { TracingUtils.recordException(span, e) - throw e - } finally { span.end() + throw e } } @@ -349,3 +356,59 @@ internal class WrappedResponseService(private val delegate: ResponseService) : R requestOptions: RequestOptions, ) = delegate.cancel(params, requestOptions) } + +/** + * Wraps a streaming Responses API stream so that completion/usage are applied to the span from the + * response.completed (or terminal) event, and the span is ended when the stream is closed. + */ +private class TracedResponseStreamResponse( + private val span: Span, + private val accumulator: ResponseAccumulator, + private val delegate: StreamResponse, +) : StreamResponse { + + @Volatile private var closed = false + + override fun stream(): Stream = + delegate.stream().map { event -> + try { + accumulator.accumulate(event) + } catch (_: IllegalStateException) { + // Terminal event already accumulated; ignore duplicate or late events + } + event + } + + override fun close() { + if (closed) return + closed = true + try { + try { + val resp = accumulator.response() + resp.usage().ifPresent { u -> + TracingUtils.setResponseAttributes( + span, + u.inputTokens().toLong(), + u.outputTokens().toLong(), + u.totalTokens().toLong(), + ) + } + resp.status().getOrNull()?.let { + span.setAttribute( + AttributeKey.stringKey("gen_ai.response.status"), + it.toString(), + ) + } + } catch (_: IllegalStateException) { + // Stream closed before a terminal event (completed/incomplete/failed) + } + } finally { + span.end() + try { + delegate.close() + } catch (_: Exception) { + // Already ended span; ensure delegate is closed best-effort + } + } + } +} diff --git a/langsmith-java-example/src/main/kotlin/com/langchain/smith/example/StreamingLangSmithExample.kt b/langsmith-java-example/src/main/kotlin/com/langchain/smith/example/StreamingLangSmithExample.kt new file mode 100644 index 00000000..04b71456 --- /dev/null +++ b/langsmith-java-example/src/main/kotlin/com/langchain/smith/example/StreamingLangSmithExample.kt @@ -0,0 +1,138 @@ +package com.langchain.smith.example + +import com.langchain.smith.wrappers.openai.OpenTelemetryConfig +import com.langchain.smith.wrappers.openai.WrappedOpenAIClient +import com.openai.models.ChatModel +import com.openai.models.chat.completions.ChatCompletionCreateParams +import kotlin.jvm.optionals.getOrNull +import java.util.concurrent.TimeUnit +import kotlin.system.exitProcess + +/** + * Example: Stream chat completions and send traces to LangSmith. + * + * - Creates a streaming chat completion request + * - Consumes the full stream (all chunks until done), accumulates content, counts chunks + * - Closes the stream so the span is ended and completion/usage are written + * - Flushes traces to LangSmith + * + * Uses the wrapped OpenAI client so streaming calls are traced and exported via OpenTelemetry. + * In LangSmith, streaming traces have span attribute `gen_ai.streaming` = true. + * + * ## Prerequisites + * Export these environment variables before running: + * - `OPENAI_API_KEY`: Your OpenAI API key + * - `LANGSMITH_API_KEY`: Your LangSmith API key + * - `LANGSMITH_PROJECT`: (optional) LangSmith project name; defaults to "default" + * - `LANGSMITH_ENDPOINT`: (optional) LangSmith API base URL; defaults to https://api.smith.langchain.com + * + * ## Running + * ```bash + * ./gradlew :langsmith-java-example:run -Pexample=StreamingLangSmith + * ``` + */ +fun main() { + println("=== Streaming Chat + LangSmith Traces Example ===\n") + + val openaiKey = System.getenv("OPENAI_API_KEY") + if (openaiKey.isNullOrEmpty()) { + System.err.println("ERROR: OPENAI_API_KEY environment variable is required!") + System.err.println(" export OPENAI_API_KEY=your_openai_api_key") + exitProcess(1) + } + + val langsmithKey = System.getenv("LANGSMITH_API_KEY") + if (langsmithKey.isNullOrEmpty()) { + System.err.println("ERROR: LANGSMITH_API_KEY environment variable is required!") + System.err.println(" export LANGSMITH_API_KEY=your_langsmith_api_key") + exitProcess(1) + } + + val projectName = System.getenv("LANGSMITH_PROJECT") ?: "default" + + val traceEndpoint = System.getenv("LANGSMITH_ENDPOINT") ?: OpenTelemetryConfig.DEFAULT_BASE_URL + println("Configuration:") + println(" LangSmith project: $projectName") + println(" Trace endpoint: $traceEndpoint/otel/v1/traces") + println() + + try { + OpenTelemetryConfig.builder() + .apiKey(langsmithKey) + .projectName(projectName) + .serviceName("streaming-langsmith-example") + .processorType(OpenTelemetryConfig.SpanProcessorType.SIMPLE) + .maxBatchSize(1) + .build() + println("✓ OpenTelemetry configured for LangSmith\n") + } catch (e: Exception) { + System.err.println("✗ Failed to configure OpenTelemetry: ${e.message}") + e.printStackTrace() + exitProcess(1) + } + + val client = WrappedOpenAIClient.fromEnv() + + + val prompt = "Count from 1 to 20" + println("Prompt: \"$prompt\"") + println("Streamed reply: ") + + var chunkCount = 0 + val fullContent = StringBuilder() + + try { + val params = ChatCompletionCreateParams.builder() + .model(ChatModel.GPT_4O_MINI) + .addUserMessage(prompt) + .maxCompletionTokens(100) + .temperature(0.0) + .build() + + // This uses the streaming API (createStreaming). Traces will have gen_ai.streaming=true. + val streamResponse = client.chat().completions().createStreaming(params) + try { + streamResponse.stream().forEach { chunk -> + chunkCount++ + val choices = chunk.choices() + if (choices.isNotEmpty()) { + val delta = choices[0].delta() + delta.content().getOrNull()?.let { content -> + fullContent.append(content) + print(content) + } + } + } + } finally { + streamResponse.close() + } + println("\n") + } catch (e: Exception) { + System.err.println("\n✗ Error during streaming: ${e.message}") + e.printStackTrace() + client.close() + exitProcess(1) + } + + client.close() + + println("Chunks received: $chunkCount") + println("Full content: \"${fullContent}\"") + println("(Trace was streaming: see gen_ai.streaming=true on the span in LangSmith)") + if (chunkCount == 0) { + System.err.println("✗ Expected at least one chunk") + } + if (fullContent.isEmpty()) { + System.err.println("✗ Expected non-empty streamed content") + } + + println("\nFlushing traces to LangSmith...") + val flushed = OpenTelemetryConfig.flush(10, TimeUnit.SECONDS) + OpenTelemetryConfig.shutdown() + if (flushed) { + println("✓ Flush completed. View traces: https://smith.langchain.com/projects/$projectName") + println(" (If you don't see traces, set LANGSMITH_DEBUG=true and run again to see export errors.)") + } else { + System.err.println("✗ Warning: Flush may not have completed. Set LANGSMITH_DEBUG=true to see export errors.") + } +}