liteai-sdk 0.3.21__py3-none-any.whl → 0.3.22__py3-none-any.whl

liteai_sdk/__init__.py

@@ -19,8 +19,7 @@ from litellm.exceptions import (
 from litellm.utils import get_valid_models
 from litellm.types.utils import LlmProviders,\
                                 ModelResponse as LiteLlmModelResponse,\
-                                ModelResponseStream as LiteLlmModelResponseStream,\
-                                Choices as LiteLlmModelResponseChoices
+                                ModelResponseStream as LiteLlmModelResponseStream
 from .debug import enable_debugging
 from .param_parser import ParamParser
 from .stream import AssistantMessageCollector
@@ -30,7 +29,7 @@ from .tool.utils import find_tool_by_name
 from .types import LlmRequestParams, GenerateTextResponse, StreamTextResponseSync, StreamTextResponseAsync
 from .types.exceptions import *
 from .types.message import ChatMessage, UserMessage, SystemMessage, AssistantMessage, ToolMessage,\
-                           MessageChunk, TextChunk, ReasoningChunk, AudioChunk, ImageChunk, ToolCallChunk,\
+                           MessageChunk, TextChunk, UsageChunk, ReasoningChunk, AudioChunk, ImageChunk, ToolCallChunk,\
                            ToolCallTuple, openai_chunk_normalizer
 from .logger import logger, enable_logging
 
@@ -146,10 +145,8 @@ class LLM:
     def generate_text_sync(self, params: LlmRequestParams) -> GenerateTextResponse:
         response = completion(**self._param_parser.parse_nonstream(params))
         response = cast(LiteLlmModelResponse, response)
-        choices = cast(list[LiteLlmModelResponseChoices], response.choices)
-        message = choices[0].message
         assistant_message = AssistantMessage\
-                            .from_litellm_message(message)\
+                            .from_litellm_message(response)\
                             .with_request_params(params)
         result: GenerateTextResponse = [assistant_message]
         if (tools_and_tool_calls := self._should_resolve_tool_calls(params, assistant_message)):
@@ -160,10 +157,8 @@ class LLM:
     async def generate_text(self, params: LlmRequestParams) -> GenerateTextResponse:
         response = await acompletion(**self._param_parser.parse_nonstream(params))
         response = cast(LiteLlmModelResponse, response)
-        choices = cast(list[LiteLlmModelResponseChoices], response.choices)
-        message = choices[0].message
         assistant_message = AssistantMessage\
-                            .from_litellm_message(message)\
+                            .from_litellm_message(response)\
                             .with_request_params(params)
         result: GenerateTextResponse = [assistant_message]
         if (tools_and_tool_calls := self._should_resolve_tool_calls(params, assistant_message)):
@@ -250,6 +245,7 @@ __all__ = [
 
     "MessageChunk",
     "TextChunk",
+    "UsageChunk",
     "ReasoningChunk",
     "AudioChunk",
     "ImageChunk",

liteai_sdk/param_parser.py

@@ -45,4 +45,5 @@ class ParamParser:
     def parse_stream(self, params: LlmRequestParams) -> ParsedParams:
         parsed = self._parse(params)
         parsed["stream"] = True
+        parsed["stream_options"] = {"include_usage": True}
         return parsed

liteai_sdk/types/message.py

@@ -6,10 +6,13 @@ from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Literal, cast
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator
 from litellm.types.utils import Message as LiteLlmMessage,\
+                                ModelResponse as LiteLlmModelResponse,\
                                 ModelResponseStream as LiteLlmModelResponseStream,\
+                                Choices as LiteLlmModelResponseChoices,\
                                 ChatCompletionAudioResponse,\
                                 ChatCompletionMessageToolCall,\
-                                ChatCompletionDeltaToolCall
+                                ChatCompletionDeltaToolCall,\
+                                Usage as LiteLlmUsage
 from litellm.types.llms.openai import (
     AllMessageValues,
     OpenAIMessageContent,
@@ -95,12 +98,16 @@ class AssistantMessage(ChatMessage):
     tool_calls: list[ChatCompletionAssistantToolCall] | None = None
     audio: ChatCompletionAudioResponse | None = None
     images: list[ChatCompletionImageURL] | None = None
+    usage: LiteLlmUsage | None = None
     role: Literal["assistant"] = "assistant"
 
     _request_params_ref: LlmRequestParams | None = PrivateAttr(default=None)
 
     @classmethod
-    def from_litellm_message(cls, message: LiteLlmMessage) -> "AssistantMessage":
+    def from_litellm_message(cls, response: LiteLlmModelResponse) -> "AssistantMessage":
+        choices = cast(list[LiteLlmModelResponseChoices], response.choices)
+        message = choices[0].message
+
         tool_calls: list[ChatCompletionAssistantToolCall] | None = None
         if (message_tool_calls := message.get("tool_calls")) is not None:
             tool_calls = [ChatCompletionAssistantToolCall(
@@ -118,6 +125,7 @@ class AssistantMessage(ChatMessage):
             tool_calls=tool_calls,
             audio=message.get("audio"),
             images=message.get("images"),
+            usage=response.get("usage"),
         )
 
     def with_request_params(self, request_params: LlmRequestParams) -> "AssistantMessage":
@@ -196,6 +204,12 @@ class SystemMessage(ChatMessage):
 class TextChunk:
     content: str
 
+@dataclasses.dataclass
+class UsageChunk:
+    input_tokens: int
+    output_tokens: int
+    total_tokens: int
+
 @dataclasses.dataclass
 class ReasoningChunk:
     content: str
@@ -239,4 +253,10 @@ def openai_chunk_normalizer(
                 tool_call.function.name,
                 tool_call.function.arguments,
                 tool_call.index))
+    if (usage := getattr(chunk, "usage", None)) is not None:
+        usage = cast(LiteLlmUsage, usage)
+        result.append(UsageChunk(
+            input_tokens=usage.prompt_tokens,
+            output_tokens=usage.completion_tokens,
+            total_tokens=usage.total_tokens))
     return result

{liteai_sdk-0.3.21.dist-info → liteai_sdk-0.3.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: liteai_sdk
-Version: 0.3.21
+Version: 0.3.22
 Summary: A wrapper of LiteLLM
 Author-email: BHznJNs <bhznjns@outlook.com>
 Requires-Python: >=3.10

{liteai_sdk-0.3.21.dist-info → liteai_sdk-0.3.22.dist-info}/RECORD

@@ -1,15 +1,15 @@
-liteai_sdk/__init__.py,sha256=KTHDeLyGGtVA-H8nJhJMBr_rKFV1h5cD5voZ8oPXI00,10608
+liteai_sdk/__init__.py,sha256=rLAihgUHA6hrh9LZSWsAoJdZWH8bIaUYAuDC1AXIb_w,10341
 liteai_sdk/debug.py,sha256=T7qIy1BeeUGlF40l9JCMMVn8pvvMJAEQeG4adQbOydA,69
 liteai_sdk/logger.py,sha256=99vJAQRKcu4CuHgZYAJ2zDQtGea6Bn3vJJrS-mtza7c,677
-liteai_sdk/param_parser.py,sha256=xykvUesZzwZNf4-n1j4JfVk0L2y_wvnSWSsHo5vjBU8,1655
+liteai_sdk/param_parser.py,sha256=KjVnTnW2cr-mJMSxhBUG__GDzTk-mKO4wlbM6Z3lODM,1714
 liteai_sdk/stream.py,sha256=T9MLmgPC8te6qvSkBOh7vkl-I4OGCKuW1kEN6RkiCe0,3176
 liteai_sdk/tool/__init__.py,sha256=c1qJaEpoYlgOCtAjFODhrSR73ZW17OuamsO__yeYAkY,12150
 liteai_sdk/tool/execute.py,sha256=1CfRlJZgqoev42fDH4vygXyEtCEEBPcRfbqaP77jxu4,2462
 liteai_sdk/tool/utils.py,sha256=Djd1-EoLPfIqgPbWWvOreozQ76NHX4FZ6OXc1evKqPM,409
 liteai_sdk/types/__init__.py,sha256=CMmweIGMgreZlbvBtRTKfvdcC7war2ApLNf-9Fz0yzc,1006
 liteai_sdk/types/exceptions.py,sha256=hIGu06htOJxfEBAHx7KTvLQr0Y8GYnBLFJFlr_IGpDs,602
-liteai_sdk/types/message.py,sha256=AnhJ5wKKcWuAt0lW3mPXpIyvUBy3u-iFLa1dpeUTp18,8785
-liteai_sdk-0.3.21.dist-info/licenses/LICENSE,sha256=cTeVgQVJJcRdm1boa2P1FBnOeXfA_egV6s4PouyrCxg,1064
-liteai_sdk-0.3.21.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-liteai_sdk-0.3.21.dist-info/METADATA,sha256=uUYWHL4MKkSTsqokLBEN2EcJd60HkinQthNXRRTabzU,3024
-liteai_sdk-0.3.21.dist-info/RECORD,,
+liteai_sdk/types/message.py,sha256=rj-h_YGdUH9x87_kToiJfyy_NV6VymZrcEnPFw1nYNU,9575
+liteai_sdk-0.3.22.dist-info/licenses/LICENSE,sha256=cTeVgQVJJcRdm1boa2P1FBnOeXfA_egV6s4PouyrCxg,1064
+liteai_sdk-0.3.22.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+liteai_sdk-0.3.22.dist-info/METADATA,sha256=Nol0OgniIOeSUQ1pOKPhnKBqRBnmJa4cmV17dDVGupk,3024
+liteai_sdk-0.3.22.dist-info/RECORD,,