feat: v2 - generations get trace level attributes (#14878)

# Pull Request Template

~~Note: merge only after https://github.com/chatwoot/ai-agents/pull/74
has been merged~~

## Description

Before:
<img width="436" height="617" alt="image"
src="https://github.com/user-attachments/assets/fd5be8dc-abab-4e01-b251-366648b998ea"
/>


After:
<img width="446" height="577" alt="image"
src="https://github.com/user-attachments/assets/8d96d2f2-c126-4cca-b3a2-f2a62b204adf"
/>


<img width="441" height="558" alt="image"
src="https://github.com/user-attachments/assets/89da4157-48a5-4fdc-9c67-9bf987e31638"
/>



## Type of change

- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide
instructions so we can reproduce. Please also list any relevant details
for your test configuration.

locally and specs

## Checklist:

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my code
- [x] I have commented on my code, particularly in hard-to-understand
areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] Any dependent changes have been merged and published in downstream
modules

---------

Co-authored-by: Sony Mathew <sony@chatwoot.com>

Gemfile

@@ -195,7 +195,7 @@ gem 'reverse_markdown'
 
 gem 'iso-639'
 gem 'ruby-openai'
-gem 'ai-agents', '>= 0.10.0'
+gem 'ai-agents', '>= 0.12.0'
 
 # TODO: Move this gem as a dependency of ai-agents
 gem 'ruby_llm', '>= 1.14.1'

Gemfile.lock

@@ -126,7 +126,7 @@ GEM
       jbuilder (~> 2)
       rails (>= 4.2, < 7.2)
       selectize-rails (~> 0.6)
-    ai-agents (0.10.0)
+    ai-agents (0.12.0)
       ruby_llm (~> 1.14)
     annotaterb (4.20.0)
       activerecord (>= 6.0.0)
@@ -1058,7 +1058,7 @@ DEPENDENCIES
   administrate (>= 0.20.1)
   administrate-field-active_storage (>= 1.0.3)
   administrate-field-belongs_to_search (>= 0.9.0)
-  ai-agents (>= 0.10.0)
+  ai-agents (>= 0.12.0)
   annotaterb
   attr_extras
   audited (~> 5.4, >= 5.4.1)

enterprise/app/services/captain/assistant/agent_runner_service.rb

@@ -155,7 +155,7 @@ class Captain::Assistant::AgentRunnerService
       span_attributes: {
         ATTR_LANGFUSE_TAGS => ['captain_v2'].to_json
       },
-      attribute_provider: ->(context_wrapper) { dynamic_trace_attributes(context_wrapper) }
+      attribute_provider: Captain::Assistant::InstrumentationAttributeProvider.new(self)
     )
     register_trace_input_callback(runner)
   end
@@ -168,7 +168,6 @@ class Captain::Assistant::AgentRunnerService
     {
       ATTR_LANGFUSE_USER_ID => state[:account_id],
       format(ATTR_LANGFUSE_METADATA, 'assistant_id') => state[:assistant_id],
-      format(ATTR_LANGFUSE_METADATA, 'conversation_id') => conversation[:id],
       format(ATTR_LANGFUSE_METADATA, 'conversation_display_id') => conversation[:display_id],
       format(ATTR_LANGFUSE_METADATA, 'channel_type') => state[:channel_type],
       format(ATTR_LANGFUSE_METADATA, 'source') => state[:source],

enterprise/app/services/captain/assistant/instrumentation_attribute_provider.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+class Captain::Assistant::InstrumentationAttributeProvider
+  include Integrations::LlmInstrumentationConstants
+
+  def initialize(service)
+    @service = service
+  end
+
+  def call(context_wrapper)
+    @service.send(:dynamic_trace_attributes, context_wrapper)
+  end
+
+  def generation_attributes(_context_wrapper, _chat, message)
+    {
+      format(ATTR_LANGFUSE_OBSERVATION_METADATA, 'generation_stage') => generation_stage(message)
+    }
+  end
+
+  private
+
+  def generation_stage(message)
+    message_has_tool_calls?(message) ? 'tool_call' : 'final_response'
+  end
+
+  def message_has_tool_calls?(message)
+    return false unless message.respond_to?(:tool_calls)
+
+    tool_calls = message.tool_calls
+    tool_calls.respond_to?(:any?) && tool_calls.any?
+  end
+end

spec/enterprise/services/captain/assistant/agent_runner_service_spec.rb

@@ -405,6 +405,47 @@ RSpec.describe Captain::Assistant::AgentRunnerService do
     end
   end
 
+  describe 'InstrumentationAttributeProvider' do
+    subject(:provider) { Captain::Assistant::InstrumentationAttributeProvider.new(service) }
+
+    let(:service) { described_class.new(assistant: assistant, conversation: conversation) }
+
+    it 'delegates root trace attributes to the service' do
+      context = {
+        state: {
+          account_id: account.id,
+          assistant_id: assistant.id,
+          conversation: { id: conversation.id, display_id: conversation.display_id }
+        }
+      }
+      context_wrapper = Struct.new(:context).new(context)
+
+      attributes = provider.call(context_wrapper)
+
+      expect(attributes).to include(
+        'langfuse.user.id' => account.id.to_s,
+        'langfuse.trace.metadata.assistant_id' => assistant.id.to_s
+      )
+    end
+
+    it 'marks final response generations for observation-level evaluators' do
+      message = instance_double(RubyLLM::Message, tool_calls: {})
+
+      attributes = provider.generation_attributes(nil, nil, message)
+
+      expect(attributes['langfuse.observation.metadata.generation_stage']).to eq('final_response')
+    end
+
+    it 'marks tool call generations separately from final responses' do
+      tool_call = instance_double(RubyLLM::ToolCall)
+      message = instance_double(RubyLLM::Message, tool_calls: { 'call_1' => tool_call })
+
+      attributes = provider.generation_attributes(nil, nil, message)
+
+      expect(attributes['langfuse.observation.metadata.generation_stage']).to eq('tool_call')
+    end
+  end
+
   describe '#build_state' do
     subject(:service) { described_class.new(assistant: assistant, conversation: conversation) }