rubyllm-observ 0.6.6 → 0.6.7

data/app/models/observ/image_generation.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Observ
+  class ImageGeneration < Observation
+    def finalize(output:, usage: {}, cost_usd: 0.0, status_message: nil)
+      merged_usage = (self.usage || {}).merge(usage.stringify_keys)
+
+      update!(
+        output: output.is_a?(String) ? output : output.to_json,
+        usage: merged_usage,
+        cost_usd: cost_usd,
+        end_time: Time.current,
+        status_message: status_message
+      )
+    end
+
+    # Image-specific helpers
+    def size
+      metadata&.dig("size")
+    end
+
+    def quality
+      metadata&.dig("quality")
+    end
+
+    def revised_prompt
+      metadata&.dig("revised_prompt")
+    end
+
+    def output_format
+      metadata&.dig("output_format") # "url" or "base64"
+    end
+
+    def mime_type
+      metadata&.dig("mime_type")
+    end
+  end
+end

data/app/models/observ/moderation.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module Observ
+  class Moderation < Observation
+    def finalize(output:, usage: {}, cost_usd: 0.0, status_message: nil)
+      merged_usage = (self.usage || {}).merge(usage.stringify_keys)
+
+      update!(
+        output: output.is_a?(String) ? output : output.to_json,
+        usage: merged_usage,
+        cost_usd: cost_usd,
+        end_time: Time.current,
+        status_message: status_message
+      )
+    end
+
+    # Moderation-specific helpers
+    def flagged?
+      metadata&.dig("flagged") || false
+    end
+
+    def categories
+      metadata&.dig("categories") || {}
+    end
+
+    def category_scores
+      metadata&.dig("category_scores") || {}
+    end
+
+    def flagged_categories
+      metadata&.dig("flagged_categories") || []
+    end
+
+    def highest_score_category
+      return nil if category_scores.empty?
+
+      category_scores.max_by { |_, score| score }&.first
+    end
+  end
+end

data/app/models/observ/null_prompt.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "mustache"
+
 module Observ
   # Null Object pattern for Prompt
   # Used when a prompt is not found, providing a fallback with the same interface
@@ -17,9 +19,36 @@ module Observ
       nil
     end
 
-    # Returns the fallback text as-is (no variable compilation)
+    # Compile prompt with Mustache templating (same as Prompt)
     def compile(variables = {})
-      @prompt
+      return @prompt if variables.empty?
+
+      Mustache.render(@prompt, variables)
+    end
+
+    # Compile with validation (raises if missing top-level variables)
+    # Note: Variables inside sections (loops) are validated at render time by Mustache
+    def compile_with_validation(variables = {})
+      required_vars = required_variables
+      provided_keys = variables.keys.map(&:to_s)
+
+      missing_vars = required_vars.reject do |var|
+        # Handle dot notation (e.g., "user.name" - check if "user" key exists)
+        root_key = var.split(".").first
+        provided_keys.include?(var) || provided_keys.include?(root_key)
+      end
+
+      if missing_vars.any?
+        raise Observ::VariableSubstitutionError, "Missing variables: #{missing_vars.join(', ')}"
+      end
+
+      compile(variables)
+    end
+
+    # Extract top-level variables from template (for validation purposes)
+    def required_variables
+      template_without_sections = strip_sections(@prompt)
+      template_without_sections.scan(/\{\{([^#\^\/!>\{\s][^}\s]*)\}\}/).flatten.uniq
     end
 
     # Null prompts are always in a "fallback" state
@@ -55,5 +84,23 @@ module Observ
     def inspect
       "#<Observ::NullPrompt name: #{name.inspect}, fallback: #{prompt[0..50].inspect}...>"
     end
+
+    private
+
+    # Strip section content from template for top-level variable extraction
+    # Removes content between {{#section}}...{{/section}} and {{^section}}...{{/section}}
+    def strip_sections(template)
+      result = template.dup
+
+      # Match sections: {{#name}}...{{/name}} or {{^name}}...{{/name}}
+      # Use non-greedy matching and handle nesting by repeating until stable
+      loop do
+        previous = result
+        result = result.gsub(/\{\{[#\^](\w+)\}\}.*?\{\{\/\1\}\}/m, "")
+        break if result == previous
+      end
+
+      result
+    end
   end
 end

data/app/models/observ/observation.rb

@@ -9,7 +9,9 @@ module Observ
 
     validates :observation_id, presence: true, uniqueness: true
     validates :start_time, presence: true
-    validates :type, presence: true, inclusion: { in: %w[Observ::Generation Observ::Span] }
+    validates :type, presence: true, inclusion: {
+    in: %w[Observ::Generation Observ::Span Observ::Embedding Observ::ImageGeneration Observ::Transcription Observ::Moderation]
+  }
 
     after_save :update_trace_metrics, if: :saved_change_to_cost_or_usage?
 

data/app/models/observ/session.rb

@@ -17,6 +17,7 @@ module Observ
 
     before_validation :set_session_id, on: :create
     before_validation :set_start_time, on: :create
+    after_commit :evaluate_guardrails, if: :just_completed?
 
     def create_trace(name: nil, input: nil, metadata: {}, tags: [])
       traces.create!(
@@ -103,6 +104,30 @@ module Observ
       instrumenter
     end
 
+    def instrument_embedding(context: {})
+      instrumenter = Observ::EmbeddingInstrumenter.new(self, context: context)
+      instrumenter.instrument!
+      instrumenter
+    end
+
+    def instrument_image_generation(context: {})
+      instrumenter = Observ::ImageGenerationInstrumenter.new(self, context: context)
+      instrumenter.instrument!
+      instrumenter
+    end
+
+    def instrument_transcription(context: {})
+      instrumenter = Observ::TranscriptionInstrumenter.new(self, context: context)
+      instrumenter.instrument!
+      instrumenter
+    end
+
+    def instrument_moderation(context: {})
+      instrumenter = Observ::ModerationInstrumenter.new(self, context: context)
+      instrumenter.instrument!
+      instrumenter
+    end
+
     def chat
       @chat ||= ::Chat.find_by(observability_session_id: session_id) if defined?(::Chat)
     end
@@ -131,5 +156,13 @@ module Observ
         ((g.end_time - g.start_time) * 1000).round(2)
       end || 0
     end
+
+    def just_completed?
+      previous_changes.key?("end_time") && end_time.present?
+    end
+
+    def evaluate_guardrails
+      GuardrailService.evaluate_session(self)
+    end
   end
 end

data/app/models/observ/trace.rb

@@ -17,6 +17,7 @@ module Observ
     validates :start_time, presence: true
 
     after_save :update_session_metrics, if: :saved_change_to_total_cost_or_total_tokens?
+    after_commit :evaluate_guardrails, if: :just_completed?
 
     def create_generation(name: "llm_generation", model: nil, metadata: {}, **options)
       observations.create!(
@@ -42,6 +43,54 @@ module Observ
       )
     end
 
+    def create_embedding(name: "embedding", model: nil, metadata: {}, **options)
+      observations.create!(
+        observation_id: SecureRandom.uuid,
+        type: "Observ::Embedding",
+        name: name,
+        model: model,
+        metadata: metadata,
+        start_time: Time.current,
+        **options.slice(:model_parameters, :parent_observation_id)
+      )
+    end
+
+    def create_image_generation(name: "image_generation", model: nil, metadata: {}, **options)
+      observations.create!(
+        observation_id: SecureRandom.uuid,
+        type: "Observ::ImageGeneration",
+        name: name,
+        model: model,
+        metadata: metadata,
+        start_time: Time.current,
+        **options.slice(:model_parameters, :parent_observation_id)
+      )
+    end
+
+    def create_transcription(name: "transcription", model: nil, metadata: {}, **options)
+      observations.create!(
+        observation_id: SecureRandom.uuid,
+        type: "Observ::Transcription",
+        name: name,
+        model: model,
+        metadata: metadata,
+        start_time: Time.current,
+        **options.slice(:model_parameters, :parent_observation_id)
+      )
+    end
+
+    def create_moderation(name: "moderation", model: nil, metadata: {}, **options)
+      observations.create!(
+        observation_id: SecureRandom.uuid,
+        type: "Observ::Moderation",
+        name: name,
+        model: model,
+        metadata: metadata,
+        start_time: Time.current,
+        **options.slice(:model_parameters, :parent_observation_id)
+      )
+    end
+
     def finalize(output: nil, metadata: {})
       merged_metadata = (self.metadata || {}).merge(metadata)
       update!(
@@ -74,13 +123,26 @@ module Observ
     end
 
     def update_aggregated_metrics
-      new_total_cost = generations.sum(:cost_usd) || 0.0
-
-      # Database-agnostic token calculation
-      new_total_tokens = generations.sum do |gen|
+      # Include generations, embeddings, image generations, transcriptions, and moderations in cost calculation
+      new_total_cost = (generations.sum(:cost_usd) || 0.0) +
+                       (embeddings.sum(:cost_usd) || 0.0) +
+                       (image_generations.sum(:cost_usd) || 0.0) +
+                       (transcriptions.sum(:cost_usd) || 0.0) +
+                       (moderations.sum(:cost_usd) || 0.0)
+
+      # Database-agnostic token calculation for generations
+      generation_tokens = generations.sum do |gen|
         gen.usage&.dig("total_tokens") || 0
       end
 
+      # Embeddings only have input tokens
+      embedding_tokens = embeddings.sum do |emb|
+        emb.usage&.dig("input_tokens") || 0
+      end
+
+      # Image generations, transcriptions, and moderations don't use tokens
+      new_total_tokens = generation_tokens + embedding_tokens
+
       update_columns(
         total_cost: new_total_cost,
         total_tokens: new_total_tokens
@@ -95,6 +157,22 @@ module Observ
       observations.where(type: "Observ::Span")
     end
 
+    def embeddings
+      observations.where(type: "Observ::Embedding")
+    end
+
+    def image_generations
+      observations.where(type: "Observ::ImageGeneration")
+    end
+
+    def transcriptions
+      observations.where(type: "Observ::Transcription")
+    end
+
+    def moderations
+      observations.where(type: "Observ::Moderation")
+    end
+
     def models_used
       generations.where.not(model: nil).distinct.pluck(:model)
     end
@@ -134,5 +212,13 @@ module Observ
     def update_session_metrics
       observ_session&.update_aggregated_metrics
     end
+
+    def just_completed?
+      previous_changes.key?("end_time") && end_time.present?
+    end
+
+    def evaluate_guardrails
+      GuardrailService.evaluate_trace(self)
+    end
   end
 end

data/app/models/observ/transcription.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Observ
+  class Transcription < Observation
+    def finalize(output:, usage: {}, cost_usd: 0.0, status_message: nil)
+      merged_usage = (self.usage || {}).merge(usage.stringify_keys)
+
+      update!(
+        output: output.is_a?(String) ? output : output.to_json,
+        usage: merged_usage,
+        cost_usd: cost_usd,
+        end_time: Time.current,
+        status_message: status_message
+      )
+    end
+
+    # Transcription-specific helpers
+    def audio_duration_s
+      metadata&.dig("audio_duration_s")
+    end
+
+    def language
+      metadata&.dig("language")
+    end
+
+    def segments_count
+      metadata&.dig("segments_count") || 0
+    end
+
+    def speakers_count
+      metadata&.dig("speakers_count")
+    end
+
+    def has_diarization?
+      metadata&.dig("has_diarization") || false
+    end
+  end
+end

data/app/services/observ/chat_instrumenter.rb

@@ -11,6 +11,7 @@ module Observ
       @current_trace = nil
       @current_tool_span = nil
       @original_ask_method = nil
+      @original_complete_method = nil
       @instrumented = false
     end
 
@@ -18,6 +19,7 @@ module Observ
       return if @instrumented
 
       wrap_ask_method
+      wrap_complete_method
       setup_event_handlers
       @instrumented = true
 
@@ -101,6 +103,66 @@ module Observ
       raise
     end
 
+    def wrap_complete_method
+      return if @original_complete_method
+
+      @original_complete_method = chat.method(:complete)
+      instrumenter = self
+
+      chat.define_singleton_method(:complete) do |**kwargs, &block|
+        instrumenter.send(:handle_complete_call, self, kwargs, block)
+      end
+    end
+
+    # Handle complete calls - similar to ask but uses existing messages
+    # instead of adding a new user message
+    def handle_complete_call(chat_instance, kwargs, block)
+      # Get the last user message for trace input
+      last_user_message = find_messages_by_role(chat_instance.messages, :user).last
+      user_message_content = last_user_message&.content
+
+      # Track if this is an ephemeral trace (created just for this call)
+      is_ephemeral_trace = @current_trace.nil?
+
+      trace = @current_trace || create_trace(
+        name: "chat.complete",
+        input: { text: user_message_content },
+        metadata: {}
+      )
+
+      model_id = extract_model_id(chat_instance)
+
+      # Extract prompt metadata from the chat's agent (if available)
+      prompt_metadata = extract_prompt_metadata(chat_instance)
+
+      generation = trace.create_generation(
+        name: "llm_call",
+        metadata: @context.merge(kwargs.slice(:temperature, :max_tokens)),
+        model: model_id,
+        model_parameters: extract_model_parameters(chat_instance),
+        **prompt_metadata
+      )
+
+      messages_snapshot = capture_messages(chat_instance)
+      generation.set_input(user_message_content, messages: messages_snapshot)
+
+      call_start_time = Time.current
+      result = @original_complete_method.call(**kwargs, &block)
+
+      finalize_generation(generation, result, call_start_time)
+
+      if is_ephemeral_trace
+        link_trace_to_message(trace, chat_instance, call_start_time)
+        trace.finalize(output: result.content)
+        @current_trace = nil
+      end
+
+      result
+    rescue StandardError => e
+      handle_error(e, trace, generation)
+      raise
+    end
+
     def setup_event_handlers
       setup_tool_call_handler
       setup_tool_result_handler
@@ -310,6 +372,26 @@ module Observ
       []
     end
 
+    # Find messages by role, handling both ActiveRecord relations and plain Arrays.
+    # ActiveRecord-backed Chat models return relations with .where(), while raw
+    # RubyLLM::Chat objects return plain Arrays.
+    #
+    # Note: We use a rescue block because some objects may claim to respond_to?(:where)
+    # but fail when the method is actually called (edge cases with proxies or custom objects).
+    def find_messages_by_role(messages, role)
+      role_str = role.to_s
+      if messages.respond_to?(:where)
+        begin
+          messages.where(role: role)
+        rescue NoMethodError
+          # Fallback to array filtering if where method doesn't actually exist
+          messages.select { |m| m.role.to_s == role_str }
+        end
+      else
+        messages.select { |m| m.role.to_s == role_str }
+      end
+    end
+
     def extract_usage(result)
       usage = {
         input_tokens: result.input_tokens || 0,
@@ -506,13 +588,21 @@ module Observ
     def link_trace_to_message(trace, chat_instance, call_start_time)
       return unless chat_instance.respond_to?(:messages)
 
-      assistant_message = chat_instance.messages
-        .where(role: "assistant")
-        .where("created_at >= ?", call_start_time)
-        .order(created_at: :desc)
-        .first
+      messages = chat_instance.messages
+      assistant_message = if messages.respond_to?(:where)
+        # ActiveRecord-backed Chat models support query methods
+        messages
+          .where(role: "assistant")
+          .where("created_at >= ?", call_start_time)
+          .order(created_at: :desc)
+          .first
+      else
+        # Raw RubyLLM::Chat objects return plain Arrays without timestamps.
+        # Get the last assistant message (most recent from this call).
+        find_messages_by_role(messages, :assistant).last
+      end
 
-      if assistant_message
+      if assistant_message&.respond_to?(:id) && assistant_message.id
         trace.update(message_id: assistant_message.id)
         Rails.logger.info "[Observability] Linked trace #{trace.trace_id} to message #{assistant_message.id}"
       end

data/app/services/observ/concerns/observable_service.rb

@@ -12,11 +12,12 @@ module Observ
     #   class MyService
     #     include Observ::Concerns::ObservableService
     #
-    #     def initialize(observability_session: nil)
+    #     def initialize(observability_session: nil, moderate: false)
     #       initialize_observability(
     #         observability_session,
     #         service_name: "my_service",
-    #         metadata: { custom: "data" }
+    #         metadata: { custom: "data" },
+    #         moderate: moderate
     #       )
     #     end
     #
@@ -24,6 +25,7 @@ module Observ
     #       with_observability do |session|
     #         # Your service logic here
     #         # Session automatically finalized on success/error
+    #         # If moderate: true, content moderation runs after finalization
     #       end
     #     end
     #   end
@@ -39,7 +41,10 @@ module Observ
       # @param session_or_false [Observ::Session, false, nil] Session to use, false to disable, nil to auto-create
       # @param service_name [String] Name of the service (used in session metadata)
       # @param metadata [Hash] Additional metadata to include in the session
-      def initialize_observability(session_or_false = nil, service_name:, metadata: {})
+      # @param moderate [Boolean] Whether to run content moderation after session finalization
+      def initialize_observability(session_or_false = nil, service_name:, metadata: {}, moderate: false)
+        @moderate_on_complete = moderate
+
         if session_or_false == false
           # Explicitly disable observability
           @observability = nil
@@ -61,6 +66,9 @@ module Observ
       # whether it succeeds or raises an error. Only sessions owned by this
       # service instance (i.e., auto-created sessions) will be finalized.
       #
+      # If moderate: true was passed to initialize_observability, content moderation
+      # will be enqueued after the session is finalized.
+      #
       # @yield [session] The observability session (may be nil if disabled)
       # @return The result of the block
       #
@@ -72,9 +80,11 @@ module Observ
       def with_observability(&block)
         result = block.call(@observability)
         finalize_service_session if @owns_session
+        enqueue_moderation if should_moderate?
         result
       rescue StandardError
         finalize_service_session if @owns_session
+        enqueue_moderation if should_moderate?
         raise
       end
 
@@ -97,6 +107,74 @@ module Observ
         @observability.instrument_chat(chat, context: context)
       end
 
+      # Instrument RubyLLM.embed for observability
+      #
+      # This wraps the RubyLLM.embed class method to automatically create traces
+      # and track embedding calls within the observability session.
+      #
+      # @param context [Hash] Additional context to include in traces
+      # @return [Observ::EmbeddingInstrumenter, nil] The instrumenter or nil if observability is disabled
+      #
+      # @example
+      #   instrument_embedding(context: { operation: "semantic_search" })
+      #   embedding = RubyLLM.embed("Search query")
+      def instrument_embedding(context: {})
+        return unless @observability
+
+        @observability.instrument_embedding(context: context)
+      end
+
+      # Instrument RubyLLM.paint for observability
+      #
+      # This wraps the RubyLLM.paint class method to automatically create traces
+      # and track image generation calls within the observability session.
+      #
+      # @param context [Hash] Additional context to include in traces
+      # @return [Observ::ImageGenerationInstrumenter, nil] The instrumenter or nil if observability is disabled
+      #
+      # @example
+      #   instrument_image_generation(context: { operation: "product_image" })
+      #   image = RubyLLM.paint("A modern logo")
+      def instrument_image_generation(context: {})
+        return unless @observability
+
+        @observability.instrument_image_generation(context: context)
+      end
+
+      # Instrument RubyLLM.transcribe for observability
+      #
+      # This wraps the RubyLLM.transcribe class method to automatically create traces
+      # and track transcription calls within the observability session.
+      #
+      # @param context [Hash] Additional context to include in traces
+      # @return [Observ::TranscriptionInstrumenter, nil] The instrumenter or nil if observability is disabled
+      #
+      # @example
+      #   instrument_transcription(context: { operation: "meeting_notes" })
+      #   transcript = RubyLLM.transcribe("meeting.wav")
+      def instrument_transcription(context: {})
+        return unless @observability
+
+        @observability.instrument_transcription(context: context)
+      end
+
+      # Instrument RubyLLM.moderate for observability
+      #
+      # This wraps the RubyLLM.moderate class method to automatically create traces
+      # and track moderation calls within the observability session.
+      #
+      # @param context [Hash] Additional context to include in traces
+      # @return [Observ::ModerationInstrumenter, nil] The instrumenter or nil if observability is disabled
+      #
+      # @example
+      #   instrument_moderation(context: { operation: "user_input_check" })
+      #   result = RubyLLM.moderate(user_input)
+      def instrument_moderation(context: {})
+        return unless @observability
+
+        @observability.instrument_moderation(context: context)
+      end
+
       private
 
       # Create a new observability session for this service
@@ -138,6 +216,33 @@ module Observ
           "[#{self.class.name}] Failed to finalize session: #{e.message}"
         )
       end
+
+      # Check if moderation should be enqueued
+      #
+      # Moderation is only enqueued when:
+      # - moderate: true was passed to initialize_observability
+      # - This service owns the session (created it)
+      # - The session exists
+      #
+      # @return [Boolean] Whether to enqueue moderation
+      def should_moderate?
+        @moderate_on_complete && @owns_session && @observability.present?
+      end
+
+      # Enqueue content moderation for the session
+      #
+      # This schedules a background job to run content moderation on all
+      # traces in the session, flagging any problematic content for review.
+      def enqueue_moderation
+        Observ::ModerationGuardrailJob.perform_later(session_id: @observability.id)
+        Rails.logger.debug(
+          "[#{self.class.name}] Moderation enqueued for session: #{@observability.session_id}"
+        )
+      rescue StandardError => e
+        Rails.logger.error(
+          "[#{self.class.name}] Failed to enqueue moderation: #{e.message}"
+        )
+      end
     end
   end
 end