Added track_model_request for litellm models

Author: psriramsnc

sygra/core/models/lite_llm/azure_openai_model.py

@@ -33,6 +33,7 @@ def __init__(self, model_config: dict[str, Any]) -> None:
     def _get_model_prefix(self) -> str:
         return "azure"
 
+    @track_model_request
     async def _generate_native_structured_output(
         self,
         input: ChatPromptValue,
@@ -68,6 +69,7 @@ async def _generate_native_structured_output(
                 api_version=self.api_version,
                 **all_params,
             )
+            self._extract_token_usage(completion)
             resp_text = completion.choices[0].model_dump()["message"]["content"]
             tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
             # Check if the request was successful based on the response status
@@ -159,6 +161,7 @@ async def _generate_text(
                 api_version=self.api_version,
                 **self.generation_params,
             )
+            self._extract_token_usage(completion)
             resp_text = completion.choices[0].model_dump()["message"]["content"]
             tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
         except openai.RateLimitError as e:

sygra/core/models/lite_llm/openai_model.py

@@ -29,6 +29,7 @@ def __init__(self, model_config: dict[str, Any]) -> None:
         self.model_config = model_config
         self.model_name = self.model_config.get("model", self.name())
 
+    @track_model_request
     async def _generate_native_structured_output(
         self,
         input: ChatPromptValue,
@@ -63,6 +64,7 @@ async def _generate_native_structured_output(
                 api_key=model_params.auth_token,
                 **all_params,
             )
+            self._extract_token_usage(completion)
             resp_text = completion.choices[0].model_dump()["message"]["content"]
             tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
             # Check if the request was successful based on the response status
@@ -153,6 +155,7 @@ async def _generate_text(
                 api_key=model_params.auth_token,
                 **self.generation_params,
             )
+            self._extract_token_usage(completion)
             resp_text = completion.choices[0].model_dump()["message"]["content"]
             tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
         except openai.RateLimitError as e:

sygra/core/models/lite_llm/vllm_model.py

@@ -33,6 +33,7 @@ def _validate_completions_api_model_support(self) -> None:
     def _get_model_prefix(self) -> str:
         return "hosted_vllm"
 
+    @track_model_request
     async def _generate_native_structured_output(
         self,
         input: ChatPromptValue,
@@ -72,6 +73,7 @@ async def _generate_native_structured_output(
                     api_key=model_params.auth_token,
                     **extra_params,
                 )
+                self._extract_token_usage(completion)
                 resp_text = completion.choices[0].model_dump()["text"]
             else:
                 # Convert input to messages
@@ -84,6 +86,7 @@ async def _generate_native_structured_output(
                     api_key=model_params.auth_token,
                     **extra_params,
                 )
+                self._extract_token_usage(completion)
                 resp_text = completion.choices[0].model_dump()["message"]["content"]
                 tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
 
@@ -157,6 +160,7 @@ async def _generate_response(
                     api_key=model_params.auth_token,
                     **self.generation_params,
                 )
+                self._extract_token_usage(completion)
                 resp_text = completion.choices[0].model_dump()["text"]
             else:
                 # Convert input to messages
@@ -169,6 +173,7 @@ async def _generate_response(
                     api_key=model_params.auth_token,
                     **self.generation_params,
                 )
+                self._extract_token_usage(completion)
                 resp_text = completion.choices[0].model_dump()["message"]["content"]
                 tool_calls = completion.choices[0].model_dump()["message"]["tool_calls"]
             # TODO: Test rate limit handling for vllm
@@ -181,7 +186,7 @@ async def _generate_response(
             logger.error(f"vLLM request failed with error: {e.message}")
             ret_code = e.status_code
         except Exception as x:
-            resp_text = f"{constants.ERROR_PREFIX} Http request failed {x}"
+            resp_text = f"{constants.ERROR_PREFIX} vLLM request failed {x}"
             logger.error(resp_text)
             rcode = self._get_status_from_body(x)
             if constants.ELEMAI_JOB_DOWN in resp_text or constants.CONNECTION_ERROR in resp_text: