address feedback

lib/llm/bedrock-command.js

@@ -14,30 +14,16 @@ class BedrockCommand {
   #modelId
 
   /**
-   * @param {object} params
-   * @param {object} params.invokeCommand The InvokeModelCommand or
+   * @param {object} input The `input` property from an InvokeModelCommand or
    * InvokeModelWithResponseStreamCommand instance that is used for the
    * conversation.
    */
-  constructor({ invokeCommand }) {
-    this.#input = invokeCommand.input
+  constructor(input) {
+    this.#input = input
     this.#body = JSON.parse(this.#input.body)
     this.#modelId = this.#input.modelId.toLowerCase()
   }
 
-  /**
-   * Given a model identifier, determine if the model is an "embedding" model.
-   *
-   * @param {string} modelId
-   *
-   * @returns {boolean}
-   */
-  static isEmbeddingModel(modelId) {
-    // At the moment, this is a simple check. If Amazon ever introduces a
-    // complex identifier, we can implement a more complicated check.
-    return modelId.toLowerCase().includes('embed')
-  }
-
   /**
    * The maximum number of tokens allowed as defined by the user.
    */
@@ -59,6 +45,15 @@ class BedrockCommand {
     return this.#modelId
   }
 
+  /**
+   * @returns {string} One of `embedding` or `completion`.
+   */
+  get modelType() {
+    // At the moment, this is a simple check. If Amazon ever introduces a
+    // complex identifier, we can implement a more complicated check.
+    return this.#modelId.toLowerCase().includes('embed') ? 'embedding' : 'completion'
+  }
+
   get prompt() {
     let result
     if (this.isTitan() === true || this.isTitanEmbed() === true) {

lib/llm/bedrock-response.js

@@ -21,6 +21,7 @@ class BedrockResponse {
   #parsedBody
   #command
   #completions = []
+  #id
 
   /**
    * @param {object} params
@@ -38,11 +39,13 @@ class BedrockResponse {
     const cmd = this.#command
     if (cmd.isAi21() === true) {
       this.#completions = this.#parsedBody.completions.map((c) => c.data.text)
+      this.#id = this.#parsedBody.id
     } else if (cmd.isClaude() === true) {
       // TODO: can we make this thing give more than one completion?
       this.#completions.push(this.#parsedBody.completion)
     } else if (cmd.isCohere() === true) {
       this.#completions = this.#parsedBody.generations.map((g) => g.text)
+      this.#id = this.#parsedBody.id
     } else if (cmd.isTitan() === true) {
       this.#completions = this.#parsedBody.results.map((r) => r.outputText)
     }
@@ -60,7 +63,7 @@ class BedrockResponse {
     } else if (cmd.isClaude() === true) {
       result = this.#parsedBody.stop_reason
     } else if (cmd.isAi21() === true) {
-      result = this.#parsedBody.results?.[0]?.finishReason.reason
+      result = this.#parsedBody.completions?.[0]?.finishReason.reason
     } else if (cmd.isCohere() === true) {
       result = this.#parsedBody.results?.generations?.[0].finish_reason
     }
@@ -71,6 +74,15 @@ class BedrockResponse {
     return this.#innerResponse.headers
   }
 
+  /**
+   * Retrieve the response identifier provided by some model responses.
+   *
+   * @returns {string|undefined}
+   */
+  get id() {
+    return this.#id
+  }
+
   get requestId() {
     return this.#innerOutput.requestId
   }

lib/llm/chat-completion-message.js

@@ -61,7 +61,7 @@ class LlmChatCompletionMessage extends LlmEvent {
     if (cmd.isTitan() === true || cmd.isClaude() === true) {
       this.id = `${this.id}-${index}`
     } else if (cmd.isAi21() === true || cmd.isCohere() === true) {
-      this.id = `${this.response.id}-${index}`
+      this.id = `${this.bedrockResponse.id}-${index}`
     }
   }
 }

lib/llm/chat-completion-summary.js

@@ -39,7 +39,7 @@ class LlmChatCompletionSummary extends LlmEvent {
     this[uct] = parseInt(this.bedrockResponse.headers['x-amzn-bedrock-output-token-count'] || 0, 10)
     this[upt] = parseInt(this.bedrockResponse.headers['x-amzn-bedrock-input-token-count'] || 0, 10)
     this[utt] = this[upt] + this[uct]
-    this[cfr] = cmd.finishReason
+    this[cfr] = this.bedrockResponse.finishReason
     this[rt] = cmd.temperature
 
     if (cmd.isTitan() === true) {

lib/llm/index.js

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+module.exports = {
+  BedrockCommand: require('./bedrock-command'),
+  BedrockResponse: require('./bedrock-response'),
+  LlmChatCompletionMessage: require('./chat-completion-message'),
+  LlmChatCompletionSummary: require('./chat-completion-summary'),
+  LlmEmbedding: require('./embedding'),
+  LlmEvent: require('./event')
+}