From 7522457740f8aabc98effc942babfa7512c7804d Mon Sep 17 00:00:00 2001 From: Aakash Bakhle <48802744+aakashb95@users.noreply.github.com> Date: Mon, 29 Jun 2026 13:50:17 +0530 Subject: [PATCH] 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: image After: image image ## 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 --- Gemfile | 2 +- Gemfile.lock | 4 +- .../captain/assistant/agent_runner_service.rb | 3 +- .../instrumentation_attribute_provider.rb | 32 +++++++++++++++ .../assistant/agent_runner_service_spec.rb | 41 +++++++++++++++++++ 5 files changed, 77 insertions(+), 5 deletions(-) create mode 100644 enterprise/app/services/captain/assistant/instrumentation_attribute_provider.rb diff --git a/Gemfile b/Gemfile index 7533cf3cf..7735dc099 100644 --- a/Gemfile +++ b/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' diff --git a/Gemfile.lock b/Gemfile.lock index 86649021f..34d8ebd38 100644 --- a/Gemfile.lock +++ b/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) diff --git a/enterprise/app/services/captain/assistant/agent_runner_service.rb b/enterprise/app/services/captain/assistant/agent_runner_service.rb index 21a9c331e..23cd6f972 100644 --- a/enterprise/app/services/captain/assistant/agent_runner_service.rb +++ b/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], diff --git a/enterprise/app/services/captain/assistant/instrumentation_attribute_provider.rb b/enterprise/app/services/captain/assistant/instrumentation_attribute_provider.rb new file mode 100644 index 000000000..b9b812b0e --- /dev/null +++ b/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 diff --git a/spec/enterprise/services/captain/assistant/agent_runner_service_spec.rb b/spec/enterprise/services/captain/assistant/agent_runner_service_spec.rb index d6e57e710..7cc72c2be 100644 --- a/spec/enterprise/services/captain/assistant/agent_runner_service_spec.rb +++ b/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) }