From dc1aee574c788094879e52378c298f265c35a311 Mon Sep 17 00:00:00 2001
From: Nir Gazit <nirga@users.noreply.github.com>
Date: Wed, 26 Jun 2024 23:50:34 +0300
Subject: [PATCH] fix(sdk): support bedrock

---
 sample-app/Gemfile                  |  2 ++
 sample-app/Gemfile.lock             | 16 ++++++++++++++-
 sample-app/bedrock.rb               | 24 +++++++++++++++++++++++
 traceloop-sdk/lib/traceloop/sdk.rb  | 30 ++++++++++++++++++++++++++++-
 traceloop-sdk/traceloop-sdk.gemspec |  2 +-
 5 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 sample-app/bedrock.rb

diff --git a/sample-app/Gemfile b/sample-app/Gemfile
index 32d757f..6c154dd 100644
--- a/sample-app/Gemfile
+++ b/sample-app/Gemfile
@@ -6,3 +6,5 @@ gem "ruby-openai"
 gem "traceloop-sdk"
 
 gem "gemini-ai", "~> 4.1"
+
+gem "aws-sdk-bedrockruntime", "~> 1.14"
diff --git a/sample-app/Gemfile.lock b/sample-app/Gemfile.lock
index c0a1d48..aae2d93 100644
--- a/sample-app/Gemfile.lock
+++ b/sample-app/Gemfile.lock
@@ -3,6 +3,18 @@ GEM
   specs:
     addressable (2.8.7)
       public_suffix (>= 2.0.2, < 7.0)
+    aws-eventstream (1.3.0)
+    aws-partitions (1.947.0)
+    aws-sdk-bedrockruntime (1.14.0)
+      aws-sdk-core (~> 3, >= 3.199.0)
+      aws-sigv4 (~> 1.1)
+    aws-sdk-core (3.199.0)
+      aws-eventstream (~> 1, >= 1.3.0)
+      aws-partitions (~> 1, >= 1.651.0)
+      aws-sigv4 (~> 1.8)
+      jmespath (~> 1, >= 1.6.1)
+    aws-sigv4 (1.8.0)
+      aws-eventstream (~> 1, >= 1.0.2)
     base64 (0.1.1)
     ethon (0.16.0)
       ffi (>= 1.15.0)
@@ -34,6 +46,7 @@ GEM
       multi_json (~> 1.11)
       os (>= 0.9, < 2.0)
       signet (>= 0.16, < 2.a)
+    jmespath (1.6.2)
     jwt (2.8.2)
       base64
     multi_json (1.15.0)
@@ -70,7 +83,7 @@ GEM
       faraday (>= 0.17.5, < 3.a)
       jwt (>= 1.5, < 3.0)
       multi_json (~> 1.10)
-    traceloop-sdk (0.0.5)
+    traceloop-sdk (0.0.8)
       opentelemetry-exporter-otlp (~> 0.26.1)
       opentelemetry-sdk (~> 1.3.1)
       opentelemetry-semantic_conventions_ai (~> 0.0.3)
@@ -81,6 +94,7 @@ PLATFORMS
   arm64-darwin-23
 
 DEPENDENCIES
+  aws-sdk-bedrockruntime (~> 1.14)
   gemini-ai (~> 4.1)
   ruby-openai
   traceloop-sdk
diff --git a/sample-app/bedrock.rb b/sample-app/bedrock.rb
new file mode 100644
index 0000000..08d95eb
--- /dev/null
+++ b/sample-app/bedrock.rb
@@ -0,0 +1,24 @@
+require 'aws-sdk-bedrockruntime'
+require "traceloop/sdk"
+
+traceloop = Traceloop::SDK::Traceloop.new
+
+model = "anthropic.claude-3-sonnet-20240229-v1:0"
+
+traceloop.llm_call(provider="bedrock", model=model) do |tracer|
+  tracer.log_prompt(user_prompt="Tell me a joke about OpenTelemetry")
+  response = Aws::BedrockRuntime::Client.new.invoke_model({
+      model_id: model,
+      content_type: "application/json",
+      accept: "*/*",
+      body: {
+        messages: [{ role: "user", content: "Tell me a joke about OpenTelemetry" }],
+        max_tokens: 4096,
+        anthropic_version: "bedrock-2023-05-31"
+      }.to_json
+    })
+  tracer.log_response(response)
+
+  body = JSON.parse(response.body.read())
+  puts body
+end
diff --git a/traceloop-sdk/lib/traceloop/sdk.rb b/traceloop-sdk/lib/traceloop/sdk.rb
index 8ae2e63..12ef4f5 100644
--- a/traceloop-sdk/lib/traceloop/sdk.rb
+++ b/traceloop-sdk/lib/traceloop/sdk.rb
@@ -46,9 +46,11 @@ def log_prompt(system_prompt="", user_prompt)
         end
 
         def log_response(response)
+          if response.respond_to?(:body)
+            log_bedrock_response(response)
           # This is Gemini specific, see -
           # https://github.com/gbaptista/gemini-ai?tab=readme-ov-file#generate_content
-          if response.has_key?("candidates")
+          elsif response.has_key?("candidates")
             log_gemini_response(response)
           else
             log_openai_response(response)
@@ -66,6 +68,32 @@ def log_gemini_response(response)
             })
         end
 
+        def log_bedrock_response(response)
+          body = JSON.parse(response.body.read())
+
+          @span.add_attributes({
+            OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_RESPONSE_MODEL => body.dig("model"),
+          })
+          if body.has_key?("usage")
+            input_tokens = body.dig("usage", "input_tokens")
+            output_tokens = body.dig("usage", "output_tokens")
+
+            @span.add_attributes({
+              OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_USAGE_TOTAL_TOKENS => input_tokens + output_tokens,
+              OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_USAGE_COMPLETION_TOKENS => output_tokens,
+              OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_USAGE_PROMPT_TOKENS => input_tokens,
+            })
+          end
+          if body.has_key?("content")
+            @span.add_attributes({
+            "#{OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_COMPLETIONS}.0.role" => body.dig("role"),
+            "#{OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_COMPLETIONS}.0.content" => body.dig("content").first.dig("text")
+            })
+          end
+
+          response.body.rewind()
+        end
+
         def log_openai_response(response)
           @span.add_attributes({
             OpenTelemetry::SemanticConventionsAi::SpanAttributes::LLM_RESPONSE_MODEL => response.dig("model"),
diff --git a/traceloop-sdk/traceloop-sdk.gemspec b/traceloop-sdk/traceloop-sdk.gemspec
index cabc8a2..8b78975 100644
--- a/traceloop-sdk/traceloop-sdk.gemspec
+++ b/traceloop-sdk/traceloop-sdk.gemspec
@@ -1,6 +1,6 @@
 Gem::Specification.new do |spec|
   spec.name        = 'traceloop-sdk'
-  spec.version     = '0.0.5'
+  spec.version     = '0.0.8'
   spec.authors     = ["Traceloop"]
   spec.email       = ['dev@traceloop.com']