Use sslr CompletionUsage in CompletionChunk message

Author: vegito22

llmstack/apps/runner/agent_controller.py

@@ -462,8 +462,8 @@ def add_llm_client_response_to_output_queue(self, response: Any):
                 AgentControllerData(
                     type=AgentControllerDataType.USAGE_DATA,
                     data=AgentUsageData(
-                        prompt_tokens=response.usage.input_tokens,
-                        completion_tokens=response.usage.output_tokens,
+                        prompt_tokens=response.usage.get_input_tokens(),
+                        completion_tokens=response.usage.get_output_tokens(),
                         total_tokens=response.usage.total_tokens,
                         source=self._provider_config.provider_config_source,
                         provider=str(self._provider_config),

llmstack/common/utils/sslr/types/chat/chat_completion_chunk.py

@@ -9,12 +9,13 @@
 from openai.types.chat.chat_completion_chunk import (  # noqa F401
     ChoiceDeltaToolCallFunction as _ChoiceDeltaToolCallFunction,
 )
-from openai.types.completion_usage import CompletionUsage
 
 from llmstack.common.utils.sslr.types.chat.chat_completion_message_param import (
     ContentPartParam,
 )
 
+from ..completions_usage import CompletionUsage
+
 
 class ChoiceDelta(_ChoiceDelta):
     content: Union[str, List[ContentPartParam]]