[FEATURE] Make observability easier

[finbarr]

Scope check

• This is core LLM communication (not application logic)
• This benefits most users (not just my use case)
• This can't be solved in application code with current RubyLLM
• I read the Contributing Guide

Due diligence

• I searched existing issues
• I checked the documentation

What problem does this solve?

It's very difficult to implement proper observability with a market leading tool like Langfuse right now.

Proposed solution

There needs to be a way to log every raw request/response.

Why this belongs in RubyLLM

This is a core use case for any serious LLM user building agentic systems. They are extremely hard to debug without proper observability.

[crmne]

Hey @finbarr if it's just about logging you can use RUBYLLM_DEBUG=true but I understood that's more than that in your PR.

[sandstrom]

You should take a look at these projects:

• https://github.com/Arize-ai/openinference
• https://github.com/traceloop/openllmetry

Maybe there are other similar projects too? But presumably some kind of standard will merge that platforms can build off.

[weilandia]

How about a middleware system that let's folks hook into the request/response cycle? This could also support caching middleware, say if we wanted to route RubyLLM.embed("Ruby is a programmer's best friend") through some cache

[weilandia]

We do this in our handrolled solution for Posthog observability:

# frozen_string_literal: true

# Middleware for adding PostHog observability to Agent::ChatClient
# Captures LLM metrics including traces, generations, token usage, latency, and errors
module Agent::ChatClient::ObservabilityMiddleware
  extend ActiveSupport::Concern
  include Agent::ObservabilityMiddleware

  included do
    prepend ObservabilityMethods
  end

  module ObservabilityMethods # rubocop:disable Metrics/ModuleLength
    def call
      return super unless observability_enabled?

      trace_id = generate_trace_id
      start_time = Time.current
      cache_hit = false
      cache_key = nil

      begin
        if cache_enabled?
          cache_key = generate_cache_key_for_observability
          cached_response = Agent::CachedChatCall.find_active(cache_key)
          cache_hit = cached_response.present?
        end

        capture_trace_start(trace_id:)

        # Execute the actual LLM call (including cache middleware if enabled)
        result = super

        if @completion&.success?
          capture_generation(trace_id:, start_time:, payload: @completion.payload, cache_hit:, cache_key:)
          capture_trace_end(trace_id:, start_time:, completion: @completion)
        else
          log_observability_skip("unsuccessful_response")
        end

        result
      rescue StandardError => e
        capture_generation_error(trace_id:, start_time:, error: e)
        capture_trace_error(trace_id:, start_time:, error: e)
        raise
      end
    end

    private

      def capture_trace_start(trace_id:)
        async_capture(
          event: "$ai_trace",
          properties: build_trace_properties(trace_id:, latency: nil, output: nil, is_error: nil)
        )
      end

      def capture_trace_end(trace_id:, start_time:, completion:)
        latency = Time.current - start_time
        output = extract_output_state(completion)

        async_capture(
          event: "$ai_trace",
          properties: build_trace_properties(trace_id:, latency:, output:, is_error: false)
        )
      end

      def capture_trace_error(trace_id:, start_time:, error:)
        latency = Time.current - start_time

        async_capture(
          event: "$ai_trace",
          properties: build_trace_properties(trace_id:, latency:, output: nil, is_error: true, error_message: error.message)
        )
      end

      def capture_generation(trace_id:, start_time:, payload:, cache_hit: false, cache_key: nil)
        latency = Time.current - start_time

        async_capture(
          event: "$ai_generation",
          properties: build_generation_properties(trace_id:, latency:, payload:, is_error: false, error: nil, cache_hit:, cache_key:)
        )
      end

      def capture_generation_error(trace_id:, start_time:, error:)
        latency = Time.current - start_time

        async_capture(
          event: "$ai_generation",
          properties: build_generation_properties(trace_id:, latency:, payload: nil, is_error: true, error:)
        )
      end

      def build_trace_properties(trace_id:, latency:, output:, is_error:, error_message: nil)
        {
          "$ai_trace_id": trace_id,
          "$ai_input_state": extract_input_state,
          "$ai_output_state": output,
          "$ai_latency": latency&.to_f,
          "$ai_span_name": extract_span_name,
          "$ai_span_id": trace_id,
          "$ai_is_error": is_error || false,
          "$ai_error": error_message,
          "$ai_client_type": self.class.name
        }.compact
      end

      def build_generation_properties(trace_id:, latency:, payload:, is_error:, error: nil, cache_hit: false, cache_key: nil) # rubocop:disable Metrics/ParameterLists
        props = base_generation_properties(trace_id:, latency:, is_error:, cache_hit:)

        add_optional_properties(props, cache_key:)
        add_metadata_properties(props)
        add_success_properties(props, payload:, cache_hit:) if payload && !is_error
        add_error_properties(props, error:) if error

        props["$ai_base_url"] = extract_base_url if respond_to?(:client)
        props.compact
      end

      def base_generation_properties(trace_id:, latency:, is_error:, cache_hit:)
        {
          "$ai_trace_id": trace_id,
          "$ai_model": model,
          "$ai_provider": extract_provider_name,
          "$ai_input": messages,
          "$ai_latency": latency.to_f,
          "$ai_is_error": is_error,
          "$ai_client_type": self.class.name,
          govly_metadata_cache_hit: cache_hit,
          govly_metadata_cache_enabled: cache_enabled?
        }
      end

      def add_optional_properties(props, cache_key:)
        props["govly_metadata_cache_key"] = cache_key if cache_key
        props["govly_metadata_prompt_class"] = prompt_instance&.class&.name if prompt_instance
        add_prompt_details(props) if prompt_instance
      end

      def add_metadata_properties(props)
        metadata = extract_metadata
        props.merge!(metadata) if metadata.present?
      end

      def add_success_properties(props, payload:, cache_hit: false)
        # For cache hits, report 0 tokens (no API cost incurred)
        if cache_hit
          props.merge!(
            "$ai_input_tokens": 0,
            "$ai_output_tokens": 0,
            "$ai_output_choices": extract_choices(payload),
            "$ai_http_status": 200
          )
        else
          input_tokens = payload.dig("usage", "prompt_tokens") || payload.dig("usage", "input_tokens")
          output_tokens = payload.dig("usage", "completion_tokens") || payload.dig("usage", "output_tokens")

          props.merge!(
            "$ai_input_tokens": input_tokens,
            "$ai_output_tokens": output_tokens,
            "$ai_output_choices": extract_choices(payload),
            "$ai_http_status": 200
          )
        end

        # No need for separate structured output field since main fields are parsed JSON
      end

      def add_error_properties(props, error:)
        props["$ai_error"] = error.message
        props["$ai_http_status"] = error_http_status(error)
      end

      def extract_input_state
        user_messages = messages.select { |m| m[:role] == "user" }
        return nil unless user_messages.any?

        user_messages.map do |msg|
          {
            role: msg[:role],
            content: [{ type: "text", text: msg[:content] }]
          }
        end
      end

      def extract_output_state(completion)
        return nil unless completion&.output_text

        [{
          role: "assistant",
          content: [{ type: "text", text: completion.output_text }]
        }]
      end

      def extract_span_name
        prompt_instance&.class&.name || self.class.name
      end

      def extract_choices(payload)
        return nil if payload&.dig("choices").blank?

        choices = payload["choices"]
        # Return as parsed array for better PostHog UI display
        choices.map do |choice|
          {
            role: choice.dig("message", "role"),
            content: choice.dig("message", "content")
          }
        end
      end

      def extract_metadata
        metadata = {}

        # First check if metadata was passed directly in parameters
        if parameters[:metadata].present?
          parameters[:metadata].each do |key, value|
            metadata["govly_metadata_#{key}"] = value
          end
        end

        # Then check request_params for metadata (from store: true tracking)
        if respond_to?(:request_params)
          params = request_params
          if params[:metadata].present?
            params[:metadata].each do |key, value|
              # Avoid duplicates by only adding if not already present
              metadata["govly_metadata_#{key}"] ||= value
            end
          end
        end

        metadata
      rescue StandardError => e
        Rails.logger.debug { "Failed to extract metadata for PostHog: #{e.message}" }
        {}
      end

      def extract_base_url
        return nil unless respond_to?(:client)

        client_config = client
        return nil unless client_config.respond_to?(:config)

        config = client_config.config
        config[:uri_base] || config[:url]
      end

      def log_observability_skip(reason)
        Rails.logger.debug { "[ObservabilityMiddleware] Skip capture: #{reason}" }
      end

      def generate_cache_key_for_observability
        # Generate the same cache key that CacheMiddleware would use
        Agent::CachedChatCall.generate_cache_key(
          prompt_class: prompt_instance&.class&.name,
          provider: provider.to_s,
          model: parameters[:model] || model,
          request: build_cache_request_for_observability
        )
      end

      def build_cache_request_for_observability
        {
          messages:,
          model: parameters[:model] || model,
          provider: provider.to_s,
          response_format: parameters[:response_format]
        }.compact
      end
  end
end

And our Cache middleware:

# frozen_string_literal: true

# Middleware for handling cache logic in Agent::EmbeddingClient
module Agent::EmbeddingClient::CacheMiddleware
  extend ActiveSupport::Concern

  included do
    prepend CacheMethods
  end

  module CacheMethods
    def embed
      return super unless cache_enabled?

      cached_response = fetch_from_cache
      if cached_response.present?
        @response = cached_response
        return input.is_a?(Array) ? embeddings : embedding
      end

      result = super

      if successful_response?
        store_to_cache(@response)
      else
        log_cache_event("skip", reason: "unsuccessful_response")
      end

      result
    rescue StandardError => e
      log_cache_event("skip", reason: "exception", error_class: e.class.name)
      raise
    end

    private

      def cache_enabled?
        # Skip caching for array inputs - unlikely to have cache hits in bulk operations
        return false if input.is_a?(Array)

        # Check if cache parameter is explicitly set, otherwise use environment default
        return @cache_enabled unless @cache_enabled.nil?

        @cache_enabled = cache.nil? ? Rails.application.config.agent_cache_default : cache
      end

      def fetch_from_cache
        cache_key_value = generate_cache_key
        cached_response = Agent::CachedEmbeddingCall.find_active(cache_key_value)

        if cached_response.present?
          log_cache_event("hit", cache_key: cache_key_value)
          cached_response
        else
          log_cache_event("miss", reason: "not_found", cache_key: cache_key_value)
          nil
        end
      end

      def store_to_cache(result)
        metadata = build_cache_metadata
        request = build_cache_request
        cache_key = generate_cache_key

        success = Agent::CachedEmbeddingCall.store(
          cache_key:,
          request:,
          response: result,
          metadata:
        )

        if success
          log_cache_event("store", provider: metadata[:provider], cache_key:)
        else
          log_cache_event("store_failed", cache_key:)
        end
      rescue StandardError => e
        Rails.logger.error "Failed to store embedding cache: #{e.message}"
      end

      def generate_cache_key
        provider_name = extract_provider_name
        Agent::CachedEmbeddingCall.generate_cache_key(
          input:,
          provider: provider_name,
          model: config.fetch(:model)
        )
      end

      def build_cache_metadata
        {
          provider: extract_provider_name,
          model: config.fetch(:model)
        }
      end

      def build_cache_request
        {
          input:,
          model: config.fetch(:model),
          provider: extract_provider_name
        }
      end

      def extract_provider_name
        provider.to_s
      end

      def successful_response?
        @response.present? &&
          @response.is_a?(Hash) &&
          @response["data"].present? &&
          @response["data"].is_a?(Array) &&
          @response["data"].all? { |item| item["embedding"].present? }
      end

      def log_cache_event(event_type, **details)
        Rails.logger.info "Agent embedding cache #{event_type}", {
          event: event_type,
          **details
        }.compact
      end
  end
end