feat: Migrate ruby llm captain (#12981)

Co-authored-by: aakashb95 <aakash@chatwoot.com>
Co-authored-by: Shivam Mishra <scm.mymail@gmail.com>

lib/integrations/llm_instrumentation.rb

@@ -1,12 +1,11 @@
 # frozen_string_literal: true
 
 require 'opentelemetry_config'
-require_relative 'llm_instrumentation_constants'
-require_relative 'llm_instrumentation_helpers'
 
 module Integrations::LlmInstrumentation
   include Integrations::LlmInstrumentationConstants
   include Integrations::LlmInstrumentationHelpers
+  include Integrations::LlmInstrumentationSpans
 
   PROVIDER_PREFIXES = {
     'openai' => %w[gpt- o1 o3 o4 text-embedding- whisper- tts-],
@@ -16,10 +15,6 @@ module Integrations::LlmInstrumentation
     'deepseek' => %w[deepseek-]
   }.freeze
 
-  def tracer
-    @tracer ||= OpentelemetryConfig.tracer
-  end
-
   def instrument_llm_call(params)
     return yield unless ChatwootApp.otel_enabled?
 
@@ -43,6 +38,7 @@ module Integrations::LlmInstrumentation
     result = nil
     executed = false
     tracer.in_span(params[:span_name]) do |span|
+      set_request_attributes(span, params)
       set_metadata_attributes(span, params)
 
       # By default, the input and output of a trace are set from the root observation
@@ -98,7 +94,12 @@ module Integrations::LlmInstrumentation
   end
 
   def record_completion(span, result)
-    set_completion_attributes(span, result) if result.is_a?(Hash)
+    if result.respond_to?(:content)
+      span.set_attribute(ATTR_GEN_AI_COMPLETION_ROLE, result.role.to_s) if result.respond_to?(:role)
+      span.set_attribute(ATTR_GEN_AI_COMPLETION_CONTENT, result.content.to_s)
+    elsif result.is_a?(Hash)
+      set_completion_attributes(span, result) if result.is_a?(Hash)
+    end
   end
 
   def set_request_attributes(span, params)
@@ -117,17 +118,4 @@ module Integrations::LlmInstrumentation
       span.set_attribute(format(ATTR_GEN_AI_PROMPT_CONTENT, idx), content.to_s)
     end
   end
-
-  def set_metadata_attributes(span, params)
-    session_id = params[:conversation_id].present? ? "#{params[:account_id]}_#{params[:conversation_id]}" : nil
-    span.set_attribute(ATTR_LANGFUSE_USER_ID, params[:account_id].to_s) if params[:account_id]
-    span.set_attribute(ATTR_LANGFUSE_SESSION_ID, session_id) if session_id.present?
-    span.set_attribute(ATTR_LANGFUSE_TAGS, [params[:feature_name]].to_json)
-
-    return unless params[:metadata].is_a?(Hash)
-
-    params[:metadata].each do |key, value|
-      span.set_attribute(format(ATTR_LANGFUSE_METADATA, key), value.to_s)
-    end
-  end
 end

lib/integrations/llm_instrumentation_helpers.rb

@@ -37,4 +37,17 @@ module Integrations::LlmInstrumentationHelpers
     span.set_attribute(ATTR_GEN_AI_RESPONSE_ERROR_CODE, error_code) if error_code
     span.status = OpenTelemetry::Trace::Status.error("API Error: #{error_code}")
   end
+
+  def set_metadata_attributes(span, params)
+    session_id = params[:conversation_id].present? ? "#{params[:account_id]}_#{params[:conversation_id]}" : nil
+    span.set_attribute(ATTR_LANGFUSE_USER_ID, params[:account_id].to_s) if params[:account_id]
+    span.set_attribute(ATTR_LANGFUSE_SESSION_ID, session_id) if session_id.present?
+    span.set_attribute(ATTR_LANGFUSE_TAGS, [params[:feature_name]].to_json)
+
+    return unless params[:metadata].is_a?(Hash)
+
+    params[:metadata].each do |key, value|
+      span.set_attribute(format(ATTR_LANGFUSE_METADATA, key), value.to_s)
+    end
+  end
 end

lib/integrations/llm_instrumentation_spans.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require 'opentelemetry_config'
+require_relative 'llm_instrumentation_constants'
+
+module Integrations::LlmInstrumentationSpans
+  include Integrations::LlmInstrumentationConstants
+
+  def tracer
+    @tracer ||= OpentelemetryConfig.tracer
+  end
+
+  def start_llm_turn_span(params)
+    return unless ChatwootApp.otel_enabled?
+
+    span = tracer.start_span(params[:span_name])
+    set_llm_turn_request_attributes(span, params)
+    set_llm_turn_prompt_attributes(span, params[:messages]) if params[:messages]
+
+    @pending_llm_turn_spans ||= []
+    @pending_llm_turn_spans.push(span)
+  rescue StandardError => e
+    Rails.logger.warn "Failed to start LLM turn span: #{e.message}"
+  end
+
+  def end_llm_turn_span(message)
+    return unless ChatwootApp.otel_enabled?
+
+    span = @pending_llm_turn_spans&.pop
+    return unless span
+
+    set_llm_turn_response_attributes(span, message) if message
+    span.finish
+  rescue StandardError => e
+    Rails.logger.warn "Failed to end LLM turn span: #{e.message}"
+  end
+
+  def start_tool_span(tool_call)
+    return unless ChatwootApp.otel_enabled?
+
+    tool_name = tool_call.name.to_s
+    span = tracer.start_span(format(TOOL_SPAN_NAME, tool_name))
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_INPUT, tool_call.arguments.to_json)
+
+    @pending_tool_spans ||= []
+    @pending_tool_spans.push(span)
+  rescue StandardError => e
+    Rails.logger.warn "Failed to start tool span: #{e.message}"
+  end
+
+  def end_tool_span(result)
+    return unless ChatwootApp.otel_enabled?
+
+    span = @pending_tool_spans&.pop
+    return unless span
+
+    output = result.is_a?(String) ? result : result.to_json
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_OUTPUT, output)
+    span.finish
+  rescue StandardError => e
+    Rails.logger.warn "Failed to end tool span: #{e.message}"
+  end
+
+  private
+
+  def set_llm_turn_request_attributes(span, params)
+    provider = determine_provider(params[:model])
+    span.set_attribute(ATTR_GEN_AI_PROVIDER, provider)
+    span.set_attribute(ATTR_GEN_AI_REQUEST_MODEL, params[:model]) if params[:model]
+    span.set_attribute(ATTR_GEN_AI_REQUEST_TEMPERATURE, params[:temperature]) if params[:temperature]
+  end
+
+  def set_llm_turn_prompt_attributes(span, messages)
+    messages.each_with_index do |msg, idx|
+      span.set_attribute(format(ATTR_GEN_AI_PROMPT_ROLE, idx), msg[:role])
+      span.set_attribute(format(ATTR_GEN_AI_PROMPT_CONTENT, idx), msg[:content])
+    end
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_INPUT, messages.to_json)
+  end
+
+  def set_llm_turn_response_attributes(span, message)
+    span.set_attribute(ATTR_GEN_AI_COMPLETION_ROLE, message.role.to_s) if message.respond_to?(:role)
+    span.set_attribute(ATTR_GEN_AI_COMPLETION_CONTENT, message.content.to_s) if message.respond_to?(:content)
+    set_llm_turn_usage_attributes(span, message)
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_OUTPUT, message.content.to_s) if message.respond_to?(:content)
+  end
+
+  def set_llm_turn_usage_attributes(span, message)
+    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, message.input_tokens) if message.respond_to?(:input_tokens) && message.input_tokens
+    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, message.output_tokens) if message.respond_to?(:output_tokens) && message.output_tokens
+  end
+end

lib/llm/config.rb

@@ -33,6 +33,7 @@ module Llm::Config
       RubyLLM.configure do |config|
         config.openai_api_key = system_api_key if system_api_key.present?
         config.openai_api_base = openai_endpoint.chomp('/') if openai_endpoint.present?
+        config.logger = Rails.logger
       end
     end