feat(bedrock):Updated the instrumentation to collect traces and metrics from imported models (#2726)

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

Author: adharshctr

packages/opentelemetry-instrumentation-bedrock/opentelemetry/instrumentation/bedrock/__init__.py

@@ -197,8 +197,13 @@ def _handle_stream_call(span, kwargs, response, metric_params):
     @dont_throw
     def stream_done(response_body):
         request_body = json.loads(kwargs.get("body"))
+        modelId = kwargs.get("modelId")
 
-        (vendor, model) = kwargs.get("modelId").split(".")
+        if modelId is not None and "." in modelId:
+            (vendor, model) = modelId.split(".")
+        else:
+            vendor = "imported_model"
+            model = kwargs.get("modelId")
 
         metric_params.vendor = vendor
         metric_params.model = model
@@ -233,6 +238,8 @@ def stream_done(response_body):
             _set_amazon_span_attributes(
                 span, request_body, response_body, metric_params
             )
+        elif vendor == "imported_model":
+            _set_imported_model_span_attributes(span, request_body, response_body, metric_params)
 
         span.end()
 
@@ -246,9 +253,13 @@ def _handle_call(span, kwargs, response, metric_params):
     )
     request_body = json.loads(kwargs.get("body"))
     response_body = json.loads(response.get("body").read())
+    modelId = kwargs.get("modelId")
 
-    (vendor, model) = kwargs.get("modelId").split(".")
-
+    if modelId is not None and "." in modelId:
+        (vendor, model) = modelId.split(".")
+    else:
+        vendor = "imported_model"
+        model = kwargs.get("modelId")
     metric_params.vendor = vendor
     metric_params.model = model
     metric_params.is_stream = False
@@ -278,6 +289,8 @@ def _handle_call(span, kwargs, response, metric_params):
         _set_llama_span_attributes(span, request_body, response_body, metric_params)
     elif vendor == "amazon":
         _set_amazon_span_attributes(span, request_body, response_body, metric_params)
+    elif vendor == "imported_model":
+        _set_imported_model_span_attributes(span, request_body, response_body, metric_params)
 
 
 def _record_usage_to_span(span, prompt_tokens, completion_tokens, metric_params):
@@ -636,6 +649,40 @@ def _set_amazon_span_attributes(span, request_body, response_body, metric_params
             )
 
 
+def _set_imported_model_span_attributes(span, request_body, response_body, metric_params):
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TYPE, LLMRequestTypeValues.COMPLETION.value
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TOP_P, request_body.get("topP")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_TEMPERATURE, request_body.get("temperature")
+    )
+    _set_span_attribute(
+        span, SpanAttributes.LLM_REQUEST_MAX_TOKENS, request_body.get("max_tokens")
+    )
+    prompt_tokens = (
+        response_body.get("usage", {}).get("prompt_tokens")
+        if response_body.get("usage", {}).get("prompt_tokens") is not None
+        else response_body.get("prompt_token_count")
+    )
+    completion_tokens = response_body.get("usage", {}).get(
+        "completion_tokens"
+    ) or response_body.get("generation_token_count")
+
+    _record_usage_to_span(span, prompt_tokens, completion_tokens, metric_params, )
+
+    if should_send_prompts():
+        _set_span_attribute(
+            span, f"{SpanAttributes.LLM_PROMPTS}.0.content", request_body.get("prompt")
+        )
+        _set_span_attribute(
+                span, f"{SpanAttributes.LLM_COMPLETIONS}.0.content",
+                response_body.get("generation"),
+            )
+
+
 def _create_metrics(meter: Meter):
     token_histogram = meter.create_histogram(
         name=Meters.LLM_TOKEN_USAGE,