Adding Google Gemini AI capabilities for AI monitoring

newrelic/config.py

@@ -2088,6 +2088,11 @@ def _process_trace_cache_import_hooks():
 
 
 def _process_module_builtin_defaults():
+    _process_module_definition(
+        "google.generativeai",
+        "newrelic.hooks.mlmodel_gemini",
+        "instrument_gemini_api_resources_chat_completion",
+    )
     _process_module_definition(
         "openai.api_resources.embedding",
         "newrelic.hooks.mlmodel_openai",

newrelic/hooks/external_botocore.py

@@ -203,6 +203,7 @@ def create_chat_completion_message_event(
 
 def extract_bedrock_titan_text_model_request(request_body, bedrock_attrs):
     request_body = json.loads(request_body)
+
     request_config = request_body.get("textGenerationConfig", {})
 
     input_message_list = [{"role": "user", "content": request_body.get("inputText")}]
@@ -225,7 +226,7 @@ def extract_bedrock_mistral_text_model_request(request_body, bedrock_attrs):
 def extract_bedrock_titan_text_model_response(response_body, bedrock_attrs):
     if response_body:
         response_body = json.loads(response_body)
-
+        
         output_message_list = [
             {"role": "assistant", "content": result["outputText"]} for result in response_body.get("results", [])
         ]
@@ -315,10 +316,10 @@ def extract_bedrock_ai21_j2_model_response(response_body, bedrock_attrs):
 
 def extract_bedrock_claude_model_request(request_body, bedrock_attrs):
     request_body = json.loads(request_body)
-
+    
     if "messages" in request_body:
         input_message_list = [
-            {"role": message.get("role", "user"), "content": message.get("content")}
+            {"role": message.get("role", "user"), "content": message.get("content")[0].get("text")}
             for message in request_body.get("messages")
         ]
     else:
@@ -333,8 +334,9 @@ def extract_bedrock_claude_model_request(request_body, bedrock_attrs):
 def extract_bedrock_claude_model_response(response_body, bedrock_attrs):
     if response_body:
         response_body = json.loads(response_body)
+
         role = response_body.get("role", "assistant")
-        content = response_body.get("content") or response_body.get("completion")
+        content = response_body.get("content")[0].get("text") or response_body.get("completion")
         output_message_list = [{"role": role, "content": content}]
         bedrock_attrs["response.choices.finish_reason"] = response_body.get("stop_reason")
         bedrock_attrs["output_message_list"] = output_message_list