braintrust 0.0.6 → 0.0.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 866cb2e797502f00cda1625ad90f4d734b4b83f0d21d8243675a933fae9df693
-  data.tar.gz: f74151b0e18b12cf19b61b1b75b2f58e784d4171f21c0996526d29c719174260
+  metadata.gz: ad055b60c4efb984bce955b1c00c684c760c849dea7a54d49a452f925dbab629
+  data.tar.gz: fe271abfac7810e53ff88efb139bc41b519799ed7dbd196ab5bb64fcbc35b62c
 SHA512:
-  metadata.gz: ad2f68a6de8d547b6a609c3393522c4ae3dfcb441a9fc841484bbbcb21de7648da7a00cd625612d98c6b99e4ad41186a2bc3fff706e17b9797e7ac514e685923
-  data.tar.gz: f0613e5fa08c07333c74467ec7830a40f72905475e35becf7a2add077168c7554046aa9a3824fe24006870338163526e8d170cfd25727af5d53416283ae03714
+  metadata.gz: a33fc58073542bf7d7dbf45092d9f3a6669d7f13fa98d3d8753d8fccd456c7e010b8a3ac0be5625e6761e164d2223f20b7fbaa3f02ead978f47398153c6c8ac2
+  data.tar.gz: b24f1377f4ec25f09c7c5e1366e1428c9164d7b3b7bdcf82331aa84d10603b250c61983ea811a13e7be1a1276558ef2995292fbefb20e99a374e59e9eaefb8b1

data/lib/braintrust/trace/contrib/github.com/crmne/ruby_llm.rb

@@ -71,15 +71,59 @@ module Braintrust
               # Check if already wrapped to make this idempotent
               return chat if chat.instance_variable_get(:@braintrust_wrapped)
 
-              # Create a wrapper module that intercepts chat.ask
+              # Create a wrapper module that intercepts chat.complete
               wrapper = create_wrapper_module(tracer_provider)
 
               # Mark as wrapped and prepend the wrapper to the chat instance
               chat.instance_variable_set(:@braintrust_wrapped, true)
               chat.singleton_class.prepend(wrapper)
+
+              # Register tool callbacks for tool span creation
+              register_tool_callbacks(chat, tracer_provider)
+
               chat
             end
 
+            # Register callbacks for tool execution tracing
+            # @param chat [RubyLLM::Chat] the chat instance
+            # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
+            def self.register_tool_callbacks(chat, tracer_provider)
+              tracer = tracer_provider.tracer("braintrust")
+
+              # Track tool spans by tool_call_id
+              tool_spans = {}
+
+              # Start tool span when tool is called
+              chat.on_tool_call do |tool_call|
+                span = tracer.start_span("ruby_llm.tool.#{tool_call.name}")
+                set_json_attr(span, "braintrust.span_attributes", {type: "tool"})
+                span.set_attribute("tool.name", tool_call.name)
+                span.set_attribute("tool.call_id", tool_call.id)
+
+                # Store tool input
+                input = {
+                  "name" => tool_call.name,
+                  "arguments" => tool_call.arguments
+                }
+                set_json_attr(span, "braintrust.input_json", input)
+
+                tool_spans[tool_call.id] = span
+              end
+
+              # End tool span when result is received
+              chat.on_tool_result do |result|
+                # Find the most recent tool span (RubyLLM doesn't pass tool_call_id to on_tool_result)
+                # The spans are processed in order, so we can use the first unfinished one
+                tool_call_id, span = tool_spans.find { |_id, s| s }
+                if span
+                  # Store tool output
+                  set_json_attr(span, "braintrust.output_json", result)
+                  span.finish
+                  tool_spans.delete(tool_call_id)
+                end
+              end
+            end
+
             # Unwrap RubyLLM to remove Braintrust tracing
             # For class-level unwrapping, removes the initialize override from the wrapper module
             # For instance-level unwrapping, clears the wrapped flag
@@ -116,50 +160,75 @@ module Braintrust
               ::RubyLLM::Chat.prepend(wrapper)
             end
 
-            # Create the wrapper module that intercepts chat.ask
+            # Create the wrapper module that intercepts chat.complete
+            # We wrap complete() instead of ask() because:
+            # - ask() internally calls complete() for the actual API call
+            # - ActiveRecord integration (acts_as_chat) calls complete() directly
+            # - This ensures all LLM calls are traced regardless of entry point
+            #
+            # Important: RubyLLM's complete() calls itself recursively for tool execution.
+            # We only create a span for the outermost call to avoid duplicate spans.
+            # Tool execution is traced separately via on_tool_call/on_tool_result callbacks.
+            #
             # @param tracer_provider [OpenTelemetry::SDK::Trace::TracerProvider] the tracer provider
             # @return [Module] the wrapper module
             def self.create_wrapper_module(tracer_provider)
               Module.new do
-                define_method(:ask) do |prompt = nil, **params, &block|
+                define_method(:complete) do |&block|
+                  # Check if we're already inside a traced complete() call
+                  # If so, just call super without creating a new span
+                  if @braintrust_in_complete
+                    if block
+                      return super(&block)
+                    else
+                      return super()
+                    end
+                  end
+
                   tracer = tracer_provider.tracer("braintrust")
 
-                  if block
-                    # Handle streaming request
-                    wrapped_block = proc do |chunk|
-                      block.call(chunk)
-                    end
-                    Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_streaming_ask(self, tracer, prompt, params, block) do |aggregated_chunks|
-                      super(prompt, **params) do |chunk|
-                        aggregated_chunks << chunk
-                        wrapped_block.call(chunk)
+                  # Mark that we're inside a complete() call
+                  @braintrust_in_complete = true
+
+                  begin
+                    if block
+                      # Handle streaming request
+                      wrapped_block = proc do |chunk|
+                        block.call(chunk)
+                      end
+                      Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_streaming_complete(self, tracer, block) do |aggregated_chunks|
+                        super(&proc do |chunk|
+                          aggregated_chunks << chunk
+                          wrapped_block.call(chunk)
+                        end)
+                      end
+                    else
+                      # Handle non-streaming request
+                      Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_non_streaming_complete(self, tracer) do
+                        super()
                       end
                     end
-                  else
-                    # Handle non-streaming request
-                    Braintrust::Trace::Contrib::Github::Crmne::RubyLLM.handle_non_streaming_ask(self, tracer, prompt, params) do
-                      super(prompt, **params)
-                    end
+                  ensure
+                    @braintrust_in_complete = false
                   end
                 end
               end
             end
 
-            # Handle streaming chat request with tracing
+            # Handle streaming complete request with tracing
             # @param chat [RubyLLM::Chat] the chat instance
             # @param tracer [OpenTelemetry::Trace::Tracer] the tracer
-            # @param prompt [String, nil] the user prompt
-            # @param params [Hash] additional parameters
             # @param block [Proc] the streaming block
-            def self.handle_streaming_ask(chat, tracer, prompt, params, block)
+            def self.handle_streaming_complete(chat, tracer, block)
               # Start span immediately for accurate timing
-              span = tracer.start_span("ruby_llm.chat.ask")
+              span = tracer.start_span("ruby_llm.chat")
 
               aggregated_chunks = []
 
               # Extract metadata and build input messages
+              # For complete(), messages are already in chat history (no prompt param)
               metadata = extract_metadata(chat, stream: true)
-              input_messages = build_input_messages(chat, prompt)
+              input_messages = build_input_messages(chat, nil)
 
               # Set input and metadata
               set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
@@ -181,19 +250,18 @@ module Braintrust
               result
             end
 
-            # Handle non-streaming chat request with tracing
+            # Handle non-streaming complete request with tracing
             # @param chat [RubyLLM::Chat] the chat instance
             # @param tracer [OpenTelemetry::Trace::Tracer] the tracer
-            # @param prompt [String, nil] the user prompt
-            # @param params [Hash] additional parameters
-            def self.handle_non_streaming_ask(chat, tracer, prompt, params)
+            def self.handle_non_streaming_complete(chat, tracer)
               # Start span immediately for accurate timing
-              span = tracer.start_span("ruby_llm.chat.ask")
+              span = tracer.start_span("ruby_llm.chat")
 
               begin
                 # Extract metadata and build input messages
+                # For complete(), messages are already in chat history (no prompt param)
                 metadata = extract_metadata(chat)
-                input_messages = build_input_messages(chat, prompt)
+                input_messages = build_input_messages(chat, nil)
                 set_json_attr(span, "braintrust.input_json", input_messages) if input_messages.any?
 
                 # Remember message count before the call (for tool call detection)
@@ -321,23 +389,62 @@ module Braintrust
             end
 
             # Build input messages array from chat history and prompt
+            # Formats messages to match OpenAI's message format
             # @param chat [RubyLLM::Chat] the chat instance
             # @param prompt [String, nil] the user prompt
             # @return [Array<Hash>] array of message hashes
             def self.build_input_messages(chat, prompt)
               input_messages = []
 
-              # Add conversation history
+              # Add conversation history, formatting each message to OpenAI format
               if chat.respond_to?(:messages) && chat.messages&.any?
-                input_messages = chat.messages.map { |m| m.respond_to?(:to_h) ? m.to_h : m }
+                input_messages = chat.messages.map { |m| format_message_for_input(m) }
               end
 
               # Add current prompt
-              input_messages << {role: "user", content: prompt} if prompt
+              input_messages << {"role" => "user", "content" => prompt} if prompt
 
               input_messages
             end
 
+            # Format a RubyLLM message to OpenAI-compatible format
+            # @param msg [Object] the RubyLLM message
+            # @return [Hash] OpenAI-formatted message
+            def self.format_message_for_input(msg)
+              formatted = {
+                "role" => msg.role.to_s
+              }
+
+              # Handle content
+              if msg.respond_to?(:content) && msg.content
+                # Convert Ruby hash notation to JSON string for tool results
+                content = msg.content
+                if msg.role.to_s == "tool" && content.is_a?(String) && content.start_with?("{:")
+                  # Ruby hash string like "{:location=>...}" - try to parse and re-serialize as JSON
+                  begin
+                    # Simple conversion: replace Ruby hash syntax with JSON
+                    content = content.gsub(/(?<=\{|, ):(\w+)=>/, '"\1":').gsub("=>", ":")
+                  rescue
+                    # Keep original if conversion fails
+                  end
+                end
+                formatted["content"] = content
+              end
+
+              # Handle tool_calls for assistant messages
+              if msg.respond_to?(:tool_calls) && msg.tool_calls&.any?
+                formatted["tool_calls"] = format_tool_calls(msg.tool_calls)
+                formatted["content"] = nil
+              end
+
+              # Handle tool_call_id for tool result messages
+              if msg.respond_to?(:tool_call_id) && msg.tool_call_id
+                formatted["tool_call_id"] = msg.tool_call_id
+              end
+
+              formatted
+            end
+
             # Capture streaming output and metrics
             # @param span [OpenTelemetry::Trace::Span] the span
             # @param aggregated_chunks [Array] the aggregated chunks
@@ -383,8 +490,9 @@ module Braintrust
               end
 
               # Check if there are tool calls in the messages history
+              # Look at messages added during this complete() call
               if chat.respond_to?(:messages) && chat.messages
-                assistant_msg = chat.messages[(messages_before_count + 1)..].find { |m|
+                assistant_msg = chat.messages[messages_before_count..].find { |m|
                   m.role.to_s == "assistant" && m.respond_to?(:tool_calls) && m.tool_calls&.any?
                 }
 

data/lib/braintrust/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Braintrust
-  VERSION = "0.0.6"
+  VERSION = "0.0.7"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: braintrust
 version: !ruby/object:Gem::Version
-  version: 0.0.6
+  version: 0.0.7
 platform: ruby
 authors:
 - Braintrust