From dfd9c3c84e106e88c9deddecb139f1edb1ee9013 Mon Sep 17 00:00:00 2001
From: Hannah Ramadan <hannahr315@gmail.com>
Date: Fri, 19 Jan 2024 14:38:20 -0800
Subject: [PATCH] Create event class files, classes

---
 lib/new_relic/agent.rb                        |  1 +
 lib/new_relic/agent/llm_event.rb              | 30 +++++++++++++
 .../agent/llm_event/chat_completion.rb        | 28 ++++++++++++
 .../llm_event/chat_completion/message.rb      | 24 +++++++++++
 .../llm_event/chat_completion/summary.rb      | 43 +++++++++++++++++++
 lib/new_relic/agent/llm_event/embedding.rb    | 36 ++++++++++++++++
 lib/new_relic/agent/llm_event/feedback.rb     | 13 ++++++
 test/new_relic/agent/llm_event_test.rb        | 19 ++++++++
 8 files changed, 194 insertions(+)
 create mode 100644 lib/new_relic/agent/llm_event.rb
 create mode 100644 lib/new_relic/agent/llm_event/chat_completion.rb
 create mode 100644 lib/new_relic/agent/llm_event/chat_completion/message.rb
 create mode 100644 lib/new_relic/agent/llm_event/chat_completion/summary.rb
 create mode 100644 lib/new_relic/agent/llm_event/embedding.rb
 create mode 100644 lib/new_relic/agent/llm_event/feedback.rb
 create mode 100644 test/new_relic/agent/llm_event_test.rb

diff --git a/lib/new_relic/agent.rb b/lib/new_relic/agent.rb
index cc93bdd18f..f3a1abb6a2 100644
--- a/lib/new_relic/agent.rb
+++ b/lib/new_relic/agent.rb
@@ -62,6 +62,7 @@ module Agent
     require 'new_relic/agent/attribute_processing'
     require 'new_relic/agent/linking_metadata'
     require 'new_relic/agent/local_log_decorator'
+    require 'new_relic/agent/llm_event'
 
     require 'new_relic/agent/instrumentation/controller_instrumentation'
 
diff --git a/lib/new_relic/agent/llm_event.rb b/lib/new_relic/agent/llm_event.rb
new file mode 100644
index 0000000000..593d960426
--- /dev/null
+++ b/lib/new_relic/agent/llm_event.rb
@@ -0,0 +1,30 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative 'llm_event/chat_completion'
+
+module NewRelic
+  module Agent
+    class LlmEvent
+
+      # response_model looks like repsonse.model
+      attr_accessor :id, :app_name, :request_id, :span_id, :transaction_id, :trace_id, :response_model, :vendor
+      INGEST_SOURCE = 'Ruby'
+
+      def initialize(id:, app_name:, request_id:, span_id:, transaction_id:, trace_id:, response_model:, vendor:, **args)
+        @id = id
+        @app_name = app_name
+        @request_id = request_id
+        @span_id = span_id
+        @transaction_id = transaction_id
+        @trace_id = trace_id
+        @response_model = response_model
+        @vendor = vendor
+      end
+
+      # Method for subclasses to override
+      def record; end
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/new_relic/agent/llm_event/chat_completion.rb b/lib/new_relic/agent/llm_event/chat_completion.rb
new file mode 100644
index 0000000000..861ab9bb50
--- /dev/null
+++ b/lib/new_relic/agent/llm_event/chat_completion.rb
@@ -0,0 +1,28 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+
+module NewRelic
+  module Agent
+    class LlmEvent
+      class ChatCompletion < NewRelic::Agent::LlmEvent
+
+        # Real metrics are written: request.max_tokens, response.number_of_messages
+        attr_accessor :api_key_last_four_digits, :conversation_id, :request_max_tokens, :response_number_of_messages
+
+        def initialize(api_key_last_four_digits:, conversation_id:, request_max_tokens:, response_number_of_messages:)
+          @api_key_last_four_digits = api_key_last_four_digits
+          @conversation_id = conversation_id
+          @request_max_tokens = request_max_tokens
+          @response_number_of_messages = response_number_of_messages
+          super
+        end
+
+      # Method for subclasses to override
+      def record; end
+      end
+    end
+  end
+end
+
diff --git a/lib/new_relic/agent/llm_event/chat_completion/message.rb b/lib/new_relic/agent/llm_event/chat_completion/message.rb
new file mode 100644
index 0000000000..95f5646db4
--- /dev/null
+++ b/lib/new_relic/agent/llm_event/chat_completion/message.rb
@@ -0,0 +1,24 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+class NewRelic
+  class Agent
+    class LlmEvent
+      class ChatCompletion
+        class Message < NewRelic::Agent::LlmEvent::ChatCompletion
+
+          attr_accessor :content, :role, :sequence, :completion_id, :is_response
+
+          def initialize
+            
+          end
+    
+          def record
+    
+          end
+        end
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/new_relic/agent/llm_event/chat_completion/summary.rb b/lib/new_relic/agent/llm_event/chat_completion/summary.rb
new file mode 100644
index 0000000000..55b6680c16
--- /dev/null
+++ b/lib/new_relic/agent/llm_event/chat_completion/summary.rb
@@ -0,0 +1,43 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+class NewRelic
+  class Agent
+    class LlmEvent
+      class ChatCompletion
+
+        attr_accessor 
+          :'request.model', 
+          :'response.organization', 
+          :'response.usage.total_tokens', 
+          :'response.usage.prompt_tokens', 
+          :'response.usage.completion_tokens',
+          :'response.choices.finish_reason',
+          :'response.headers.llmVersion',
+          :'response.headers.ratelimitLimitRequests',
+          :'response.headers.ratelimitLimitTokens',
+          :'response.headers.ratelimitResetTokens',
+          :'response.headers.ratelimitResetRequests',
+          :'response.headers.ratelimitRemainingTokens',
+          :'response.headers.ratelimitRemainingRequests',
+          :duration,
+          :'request.temperature',
+          :error
+
+
+
+        class Summary < NewRelic::Agent::LlmEvent::ChatCompletion
+          
+          def initialize
+            
+          end
+    
+          def record
+    
+          end
+        end
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/new_relic/agent/llm_event/embedding.rb b/lib/new_relic/agent/llm_event/embedding.rb
new file mode 100644
index 0000000000..f0123050e4
--- /dev/null
+++ b/lib/new_relic/agent/llm_event/embedding.rb
@@ -0,0 +1,36 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+class NewRelic
+  class Agent
+    class LlmEvent
+      class Embedding < NewRelic::Agent::LlmEvent
+
+        attr_accessor 
+          :input, 
+          :'request.model', 
+          :'response.organization', 
+          :'response.usage.total_tokens', 
+          :'response.usage.prompt_tokens',
+          :'response.headers.llmVersion',
+          :'response.headers.ratelimitLimitRequests'
+          :'response.headers.ratelimitLimitTokens',
+          :'response.headers.ratelimitResetTokens',
+          :'response.headers.ratelimitResetRequests',
+          :'response.headers.ratelimitRemainingTokens',
+          :'response.headers.ratelimitRemainingRequests',
+          :duration,
+          :error
+
+        def initialize
+
+        end
+
+        def record
+
+        end
+      end
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/new_relic/agent/llm_event/feedback.rb b/lib/new_relic/agent/llm_event/feedback.rb
new file mode 100644
index 0000000000..c483de8d37
--- /dev/null
+++ b/lib/new_relic/agent/llm_event/feedback.rb
@@ -0,0 +1,13 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+
+class NewRelic
+  class Agent
+    class LlmEvent
+      class Feedback < NewRelic::Agent::LlmEvent
+      end
+    end
+  end
+end
diff --git a/test/new_relic/agent/llm_event_test.rb b/test/new_relic/agent/llm_event_test.rb
new file mode 100644
index 0000000000..924cb95503
--- /dev/null
+++ b/test/new_relic/agent/llm_event_test.rb
@@ -0,0 +1,19 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative '../../test_helper'
+
+
+module NewRelic::Agent
+  class LlmEventTest < Minitest::Test
+    # def test_attributes
+    #   NewRelic::Agent::LlmEvent.new
+
+    # end
+
+    def test_attributes_chat
+      NewRelic::Agent::LlmEvent::ChatCompletion.new
+    end
+  end
+end
\ No newline at end of file