tracekit 0.2.2 → 0.2.4

data/lib/tracekit/llm/anthropic_instrumentation.rb

@@ -0,0 +1,218 @@
+# frozen_string_literal: true
+
+require_relative "common"
+
+module Tracekit
+  module LLM
+    module AnthropicInstrumentation
+      module_function
+
+      def install(tracer)
+        begin
+          require "anthropic"
+        rescue LoadError
+          # anthropic gem not available, check if it's already defined (e.g. in tests)
+          return false unless defined?(::Anthropic::Client)
+        end
+
+        return false unless defined?(::Anthropic::Client)
+
+        instrumentation_mod = Module.new do
+          define_method(:messages) do |**params|
+            # When called with no parameters, return the Messages::Client (for batches etc.)
+            return super(**params) unless params[:parameters]
+
+            parameters = params[:parameters]
+            model = parameters[:model] || parameters["model"] || "unknown"
+            stream_proc = parameters[:stream] || parameters["stream"]
+            is_streaming = stream_proc.is_a?(Proc)
+            capture = Common.capture_content?
+
+            span = tracer.start_span("chat #{model}", kind: :client)
+
+            begin
+              Common.set_request_attributes(span,
+                provider: "anthropic",
+                model: model,
+                max_tokens: parameters[:max_tokens] || parameters["max_tokens"],
+                temperature: parameters[:temperature] || parameters["temperature"],
+                top_p: parameters[:top_p] || parameters["top_p"]
+              )
+
+              # Capture input content
+              if capture
+                system_prompt = parameters[:system] || parameters["system"]
+                Common.capture_system_instructions(span, system_prompt) if system_prompt
+                messages = parameters[:messages] || parameters["messages"]
+                Common.capture_input_messages(span, messages) if messages
+              end
+
+              if is_streaming
+                # Wrap the user's stream proc to accumulate span data
+                accumulator = AnthropicStreamAccumulator.new(span, capture)
+                wrapper_proc = proc do |event|
+                  accumulator.process_event(event)
+                  stream_proc.call(event)
+                end
+
+                # Replace stream proc with our wrapper
+                wrapped_params = parameters.merge(stream: wrapper_proc)
+                result = super(parameters: wrapped_params)
+                accumulator.finalize
+                result
+              else
+                result = super(**params)
+                handle_anthropic_response(span, result, capture)
+                result
+              end
+            rescue => e
+              Common.set_error_attributes(span, e)
+              span.finish
+              raise
+            end
+          end
+
+          private
+
+          def handle_anthropic_response(span, result, capture)
+            # Anthropic response: { id, type, role, content, model, stop_reason, usage }
+            content_blocks = result["content"] || result[:content] || []
+            usage = result["usage"] || result[:usage] || {}
+
+            Common.set_response_attributes(span,
+              model: result["model"] || result[:model],
+              id: result["id"] || result[:id],
+              finish_reasons: [(result["stop_reason"] || result[:stop_reason])].compact,
+              input_tokens: usage["input_tokens"] || usage[:input_tokens],
+              output_tokens: usage["output_tokens"] || usage[:output_tokens]
+            )
+
+            # Cache tokens (Anthropic-specific)
+            cache_creation = usage["cache_creation_input_tokens"] || usage[:cache_creation_input_tokens]
+            cache_read = usage["cache_read_input_tokens"] || usage[:cache_read_input_tokens]
+            span.set_attribute("gen_ai.usage.cache_creation.input_tokens", cache_creation) if cache_creation
+            span.set_attribute("gen_ai.usage.cache_read.input_tokens", cache_read) if cache_read
+
+            # Tool calls from content blocks
+            content_blocks.each do |block|
+              block_type = block["type"] || block[:type]
+              if block_type == "tool_use"
+                input_val = block["input"] || block[:input]
+                args = input_val.is_a?(String) ? input_val : JSON.generate(input_val)
+                Common.record_tool_call(span,
+                  name: block["name"] || block[:name] || "unknown",
+                  id: block["id"] || block[:id],
+                  arguments: args
+                )
+              end
+            end
+
+            # Output content capture
+            if capture && content_blocks.any?
+              Common.capture_output_messages(span, content_blocks)
+            end
+          rescue => _e
+            # Never break user code
+          ensure
+            span.finish
+          end
+        end
+
+        ::Anthropic::Client.prepend(instrumentation_mod)
+        true
+      end
+
+      # Accumulates streaming event data for span attributes
+      class AnthropicStreamAccumulator
+        def initialize(span, capture_content)
+          @span = span
+          @capture = capture_content
+          @model = nil
+          @id = nil
+          @stop_reason = nil
+          @input_tokens = nil
+          @output_tokens = nil
+          @cache_creation_tokens = nil
+          @cache_read_tokens = nil
+          @output_chunks = []
+          @tool_calls = {}
+          @current_block_index = 0
+        end
+
+        def process_event(event)
+          event_type = event["type"] || event[:type]
+
+          case event_type
+          when "message_start"
+            message = event["message"] || event[:message] || {}
+            @model = message["model"] || message[:model]
+            @id = message["id"] || message[:id]
+            usage = message["usage"] || message[:usage] || {}
+            @input_tokens = usage["input_tokens"] || usage[:input_tokens]
+            @cache_creation_tokens = usage["cache_creation_input_tokens"] || usage[:cache_creation_input_tokens]
+            @cache_read_tokens = usage["cache_read_input_tokens"] || usage[:cache_read_input_tokens]
+
+          when "content_block_start"
+            @current_block_index = event["index"] || event[:index] || @current_block_index
+            cb = event["content_block"] || event[:content_block] || {}
+            if (cb["type"] || cb[:type]) == "tool_use"
+              @tool_calls[@current_block_index] = {
+                name: cb["name"] || cb[:name] || "unknown",
+                id: cb["id"] || cb[:id],
+                arguments: ""
+              }
+            end
+
+          when "content_block_delta"
+            delta = event["delta"] || event[:delta] || {}
+            delta_type = delta["type"] || delta[:type]
+            if delta_type == "text_delta" && @capture
+              text = delta["text"] || delta[:text]
+              @output_chunks << text if text
+            elsif delta_type == "input_json_delta"
+              partial = delta["partial_json"] || delta[:partial_json]
+              idx = event["index"] || event[:index] || @current_block_index
+              if partial && @tool_calls[idx]
+                @tool_calls[idx][:arguments] += partial
+              end
+            end
+
+          when "message_delta"
+            delta = event["delta"] || event[:delta] || {}
+            @stop_reason = delta["stop_reason"] || delta[:stop_reason] if delta["stop_reason"] || delta[:stop_reason]
+            usage = event["usage"] || event[:usage] || {}
+            @output_tokens = usage["output_tokens"] || usage[:output_tokens] if usage["output_tokens"] || usage[:output_tokens]
+          end
+        rescue => _e
+          # Never fail on event processing
+        end
+
+        def finalize
+          Common.set_response_attributes(@span,
+            model: @model,
+            id: @id,
+            finish_reasons: @stop_reason ? [@stop_reason] : nil,
+            input_tokens: @input_tokens,
+            output_tokens: @output_tokens
+          )
+
+          @span.set_attribute("gen_ai.usage.cache_creation.input_tokens", @cache_creation_tokens) if @cache_creation_tokens
+          @span.set_attribute("gen_ai.usage.cache_read.input_tokens", @cache_read_tokens) if @cache_read_tokens
+
+          @tool_calls.each_value do |tc|
+            Common.record_tool_call(@span, **tc)
+          end
+
+          if @capture && @output_chunks.any?
+            full_content = @output_chunks.join
+            Common.capture_output_messages(@span, [{ "type" => "text", "text" => full_content }])
+          end
+        rescue => _e
+          # Never break user code
+        ensure
+          @span.finish
+        end
+      end
+    end
+  end
+end

data/lib/tracekit/llm/common.rb

@@ -0,0 +1,118 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Tracekit
+  module LLM
+    module Common
+      # Pattern-based PII regexes (all replaced with plain [REDACTED])
+      SENSITIVE_KEY_PATTERN = /\A(password|passwd|pwd|secret|token|key|credential|api_key|apikey)\z/i
+      EMAIL_PATTERN = /[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}/
+      SSN_PATTERN = /\b\d{3}-\d{2}-\d{4}\b/
+      CREDIT_CARD_PATTERN = /\b\d{4}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}\b/
+      AWS_KEY_PATTERN = /\bAKIA[0-9A-Z]{16}\b/
+      BEARER_PATTERN = /Bearer\s+[A-Za-z0-9\-._~+\/]+=*/
+      STRIPE_PATTERN = /\bsk_live_[a-zA-Z0-9]+/
+      JWT_PATTERN = /\beyJ[A-Za-z0-9\-_]+\.eyJ[A-Za-z0-9\-_]+\.[A-Za-z0-9\-_]+/
+      PRIVATE_KEY_PATTERN = /-----BEGIN\s+(?:RSA\s+)?PRIVATE\s+KEY-----/
+
+      CONTENT_PATTERNS = [
+        EMAIL_PATTERN, SSN_PATTERN, CREDIT_CARD_PATTERN, AWS_KEY_PATTERN,
+        BEARER_PATTERN, STRIPE_PATTERN, JWT_PATTERN, PRIVATE_KEY_PATTERN
+      ].freeze
+
+      module_function
+
+      def scrub_pii(content)
+        # Try JSON key-based scrubbing first
+        begin
+          parsed = JSON.parse(content)
+          scrubbed = scrub_object(parsed)
+          return JSON.generate(scrubbed)
+        rescue JSON::ParserError
+          # Not JSON, fall through to pattern scrubbing
+        end
+        scrub_patterns(content)
+      end
+
+      def scrub_patterns(str)
+        result = str.dup
+        CONTENT_PATTERNS.each { |pat| result.gsub!(pat, "[REDACTED]") }
+        result
+      end
+
+      def scrub_object(obj)
+        case obj
+        when Hash
+          obj.each_with_object({}) do |(k, v), h|
+            if SENSITIVE_KEY_PATTERN.match?(k.to_s)
+              h[k] = "[REDACTED]"
+            else
+              h[k] = scrub_object(v)
+            end
+          end
+        when Array
+          obj.map { |item| scrub_object(item) }
+        when String
+          scrub_patterns(obj)
+        else
+          obj
+        end
+      end
+
+      def capture_content?
+        env_val = ENV["TRACEKIT_LLM_CAPTURE_CONTENT"]
+        return env_val.downcase == "true" || env_val == "1" if env_val
+        false
+      end
+
+      def set_request_attributes(span, provider:, model:, max_tokens: nil, temperature: nil, top_p: nil)
+        span.set_attribute("gen_ai.operation.name", "chat")
+        span.set_attribute("gen_ai.system", provider)
+        span.set_attribute("gen_ai.request.model", model)
+        span.set_attribute("gen_ai.request.max_tokens", max_tokens) if max_tokens
+        span.set_attribute("gen_ai.request.temperature", temperature) if temperature
+        span.set_attribute("gen_ai.request.top_p", top_p) if top_p
+      end
+
+      def set_response_attributes(span, model: nil, id: nil, finish_reasons: nil, input_tokens: nil, output_tokens: nil)
+        span.set_attribute("gen_ai.response.model", model) if model
+        span.set_attribute("gen_ai.response.id", id) if id
+        span.set_attribute("gen_ai.response.finish_reasons", finish_reasons) if finish_reasons&.any?
+        span.set_attribute("gen_ai.usage.input_tokens", input_tokens) if input_tokens
+        span.set_attribute("gen_ai.usage.output_tokens", output_tokens) if output_tokens
+      end
+
+      def set_error_attributes(span, error)
+        span.set_attribute("error.type", error.class.name)
+        span.status = OpenTelemetry::Trace::Status.error(error.message)
+        span.record_exception(error)
+      end
+
+      def record_tool_call(span, name:, id: nil, arguments: nil)
+        attrs = { "gen_ai.tool.name" => name }
+        attrs["gen_ai.tool.call.id"] = id if id
+        attrs["gen_ai.tool.call.arguments"] = arguments if arguments
+        span.add_event("gen_ai.tool.call", attributes: attrs)
+      end
+
+      def capture_input_messages(span, messages)
+        return unless messages
+        serialized = JSON.generate(messages)
+        span.set_attribute("gen_ai.input.messages", scrub_pii(serialized))
+      end
+
+      def capture_output_messages(span, content)
+        return unless content
+        serialized = JSON.generate(content)
+        span.set_attribute("gen_ai.output.messages", scrub_pii(serialized))
+      end
+
+      def capture_system_instructions(span, system)
+        return unless system
+        serialized = system.is_a?(String) ? system : JSON.generate(system)
+        span.set_attribute("gen_ai.system_instructions", scrub_pii(serialized))
+      end
+    end
+  end
+end

data/lib/tracekit/llm/openai_instrumentation.rb

@@ -0,0 +1,201 @@
+# frozen_string_literal: true
+
+require_relative "common"
+
+module Tracekit
+  module LLM
+    module OpenAIInstrumentation
+      module_function
+
+      def install(tracer)
+        # Try to load the OpenAI gem
+        begin
+          require "openai"
+        rescue LoadError
+          # openai gem not available, check if it's already defined (e.g. in tests)
+          return false unless defined?(::OpenAI::Client)
+        end
+
+        client_class = ::OpenAI::Client
+        return false unless client_class
+
+        # Create the prepend module dynamically with tracer closure
+        instrumentation_mod = Module.new do
+          define_method(:chat) do |parameters: {}|
+            model = parameters[:model] || parameters["model"] || "unknown"
+            stream_proc = parameters[:stream] || parameters["stream"]
+            is_streaming = stream_proc.is_a?(Proc)
+            capture = Common.capture_content?
+
+            span = tracer.start_span("chat #{model}", kind: :client)
+
+            begin
+              Common.set_request_attributes(span,
+                provider: "openai",
+                model: model,
+                max_tokens: parameters[:max_tokens] || parameters["max_tokens"] || parameters[:max_completion_tokens] || parameters["max_completion_tokens"],
+                temperature: parameters[:temperature] || parameters["temperature"],
+                top_p: parameters[:top_p] || parameters["top_p"]
+              )
+
+              # Capture input content
+              if capture
+                messages = parameters[:messages] || parameters["messages"]
+                if messages
+                  system_msgs = messages.select { |m| (m[:role] || m["role"]) == "system" }
+                  non_system = messages.reject { |m| (m[:role] || m["role"]) == "system" }
+                  Common.capture_system_instructions(span, system_msgs) if system_msgs.any?
+                  Common.capture_input_messages(span, non_system)
+                end
+              end
+
+              if is_streaming
+                # ruby-openai handles streaming via proc callback internally.
+                # The chat method returns the final response hash, not an enumerator.
+                # We wrap the user's proc to accumulate span data from each chunk.
+                accumulator = OpenAIStreamAccumulator.new(span, capture)
+                wrapper_proc = proc do |chunk, bytesize|
+                  accumulator.process_chunk(chunk)
+                  # Call original proc with same args
+                  if stream_proc.arity == 2 || stream_proc.arity < 0
+                    stream_proc.call(chunk, bytesize)
+                  else
+                    stream_proc.call(chunk)
+                  end
+                end
+
+                # Inject stream_options.include_usage for token counting
+                params = parameters.dup
+                so = params[:stream_options] || params["stream_options"] || {}
+                unless so[:include_usage] || so["include_usage"]
+                  params[:stream_options] = so.merge(include_usage: true)
+                end
+                params[:stream] = wrapper_proc
+
+                result = super(parameters: params)
+                accumulator.finalize
+                result
+              else
+                result = super(parameters: parameters)
+
+                # Non-streaming response handling
+                handle_response(span, result, capture)
+                result
+              end
+            rescue => e
+              Common.set_error_attributes(span, e)
+              span.finish
+              raise
+            end
+          end
+
+          private
+
+          def handle_response(span, result, capture)
+            choices = result.dig("choices") || []
+            Common.set_response_attributes(span,
+              model: result["model"],
+              id: result["id"],
+              finish_reasons: choices.map { |c| c["finish_reason"] }.compact,
+              input_tokens: result.dig("usage", "prompt_tokens"),
+              output_tokens: result.dig("usage", "completion_tokens")
+            )
+
+            # Tool calls
+            choices.each do |choice|
+              (choice.dig("message", "tool_calls") || []).each do |tc|
+                Common.record_tool_call(span,
+                  name: tc.dig("function", "name") || "unknown",
+                  id: tc["id"],
+                  arguments: tc.dig("function", "arguments")
+                )
+              end
+            end
+
+            # Output content capture
+            if capture && choices.any?
+              output_msgs = choices.map { |c| c["message"] }.compact
+              Common.capture_output_messages(span, output_msgs) if output_msgs.any?
+            end
+          rescue => _e
+            # Never break user code
+          ensure
+            span.finish
+          end
+        end
+
+        client_class.prepend(instrumentation_mod)
+        true
+      end
+
+      # Accumulates streaming chunk data for span attributes via proc interception
+      class OpenAIStreamAccumulator
+        def initialize(span, capture_content)
+          @span = span
+          @capture = capture_content
+          @model = nil
+          @id = nil
+          @finish_reason = nil
+          @input_tokens = nil
+          @output_tokens = nil
+          @output_chunks = []
+          @tool_calls = {}
+        end
+
+        def process_chunk(chunk)
+          @model ||= chunk.dig("model")
+          @id ||= chunk.dig("id")
+
+          if (usage = chunk["usage"])
+            @input_tokens = usage["prompt_tokens"] if usage["prompt_tokens"]
+            @output_tokens = usage["completion_tokens"] if usage["completion_tokens"]
+          end
+
+          (chunk["choices"] || []).each do |choice|
+            @finish_reason = choice["finish_reason"] if choice["finish_reason"]
+            delta = choice["delta"] || {}
+            @output_chunks << delta["content"] if @capture && delta["content"]
+
+            (delta["tool_calls"] || []).each do |tc|
+              idx = tc["index"] || 0
+              if @tool_calls[idx]
+                @tool_calls[idx][:arguments] = (@tool_calls[idx][:arguments] || "") + (tc.dig("function", "arguments") || "")
+              else
+                @tool_calls[idx] = {
+                  name: tc.dig("function", "name") || "unknown",
+                  id: tc["id"],
+                  arguments: tc.dig("function", "arguments") || ""
+                }
+              end
+            end
+          end
+        rescue => _e
+          # Never fail on chunk processing
+        end
+
+        def finalize
+          Common.set_response_attributes(@span,
+            model: @model,
+            id: @id,
+            finish_reasons: @finish_reason ? [@finish_reason] : nil,
+            input_tokens: @input_tokens,
+            output_tokens: @output_tokens
+          )
+
+          @tool_calls.each_value do |tc|
+            Common.record_tool_call(@span, **tc)
+          end
+
+          if @capture && @output_chunks.any?
+            full_content = @output_chunks.join
+            Common.capture_output_messages(@span, [{ "role" => "assistant", "content" => full_content }])
+          end
+        rescue => _e
+          # Never break user code
+        ensure
+          @span.finish
+        end
+      end
+    end
+  end
+end

data/lib/tracekit/sdk.rb

@@ -90,6 +90,9 @@ module Tracekit
       # Initialize OpenTelemetry tracer
       setup_tracing(traces_endpoint)
 
+      # Initialize LLM instrumentation (auto-detect providers)
+      setup_llm_instrumentation if defined?(Tracekit::LLM)
+
       # Initialize metrics registry
       @metrics_registry = Metrics::Registry.new(metrics_endpoint, config.api_key, config.service_name)
 
@@ -152,6 +155,32 @@ module Tracekit
 
     private
 
+    def setup_llm_instrumentation
+      llm_config = @config.llm || {}
+      return unless llm_config.fetch(:enabled, true)
+
+      tracer = OpenTelemetry.tracer_provider.tracer("tracekit-llm", Tracekit::VERSION)
+
+      # Set capture_content env var from config if not already set
+      if llm_config[:capture_content] && !ENV.key?("TRACEKIT_LLM_CAPTURE_CONTENT")
+        ENV["TRACEKIT_LLM_CAPTURE_CONTENT"] = "true"
+      end
+
+      if llm_config.fetch(:openai, true)
+        if Tracekit::LLM::OpenAIInstrumentation.install(tracer)
+          puts "TraceKit: OpenAI LLM instrumentation enabled"
+        end
+      end
+
+      if llm_config.fetch(:anthropic, true)
+        if Tracekit::LLM::AnthropicInstrumentation.install(tracer)
+          puts "TraceKit: Anthropic LLM instrumentation enabled"
+        end
+      end
+    rescue => e
+      puts "TraceKit: LLM instrumentation setup failed: #{e.message}"
+    end
+
     def setup_tracing(traces_endpoint)
       OpenTelemetry::SDK.configure do |c|
         c.service_name = @config.service_name