changed the metering data to take model_id

Author: EC2 Default User

lib/idp_common_pkg/idp_common/bedrock/client.py

@@ -296,10 +296,7 @@ def invoke_model(
                 inference_config["maxTokens"] = max_tokens
 
         # Add additional model fields if needed
-        additional_model_fields = {}
-        
-        # Initialize use_model_id early to avoid UnboundLocalError
-        use_model_id = model_id
+        additional_model_fields = {}     
 
         # Handle top_k parameter
         if top_k is not None:
@@ -331,6 +328,7 @@ def invoke_model(
                 additional_model_fields["inferenceConfig"]["topK"] = int(top_k)
 
         # Add 1M context headers if needed
+        use_model_id = model_id
         if model_id and model_id.endswith(':1m'):
             use_model_id = model_id[:-3]  # Remove ':1m'
             if additional_model_fields is None:
@@ -458,7 +456,7 @@ def _invoke_with_retry(
             response_with_metering = {
                 "response": response,
                 "metering": {
-                    f"{context}/bedrock/{'modelId'}": {
+                    f"{context}/bedrock/{model_id}": {
                         **usage
                     }
                 }