Merge pull request #12609 from AnuGayan/master-AI-API-temp

Fix issues in AI Subscription throttling policy deployment
wso2 · Sep 27, 2024 · e35fff6 · e35fff6
2 parents 2742c8a + c6c5144
commit e35fff6
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/...va/org/wso2/carbon/apimgt/gateway/throttling/publisher/DataProcessAndPublishingAgent.java b/...va/org/wso2/carbon/apimgt/gateway/throttling/publisher/DataProcessAndPublishingAgent.java
@@ -218,9 +218,18 @@ public void setDataReference(String applicationLevelThrottleKey, String applicat
             Map<String, String> responseMetadata = (Map<String, String>) ((Axis2MessageContext) messageContext)
                     .getAxis2MessageContext().getProperty(AI_API_RESPONSE_METADATA);
             if (responseMetadata != null) {
-                totalTokens = Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_TOTAL_TOKEN_COUNT));
-                promptTokens = Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_PROMPT_TOKEN_COUNT));
-                completionTokens = Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_COMPLETION_TOKEN_COUNT));
+                if (null != responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_TOTAL_TOKEN_COUNT)) {
+                    totalTokens =
+                            Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_TOTAL_TOKEN_COUNT));
+                }
+                if (null != responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_PROMPT_TOKEN_COUNT)) {
+                    promptTokens =
+                            Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_PROMPT_TOKEN_COUNT));
+                }
+                if (null != responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_COMPLETION_TOKEN_COUNT)) {
+                    completionTokens =
+                            Long.parseLong(responseMetadata.get(LLM_PROVIDER_SERVICE_METADATA_COMPLETION_TOKEN_COUNT));
+                }
             }
         }
     }