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

liteai_sdk/__init__.py

@@ -3,7 +3,7 @@ import queue
 from typing import cast
 from collections.abc import AsyncGenerator, Generator
 from litellm import CustomStreamWrapper, completion, acompletion
-from litellm.exceptions import (  
+from litellm.exceptions import (
     AuthenticationError,
     PermissionDeniedError,
     RateLimitError,
@@ -19,18 +19,18 @@ 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
-from .tool import ToolFn, ToolDef, RawToolDef, ToolLike
 from .tool.execute import execute_tool_sync, execute_tool
+from .tool.toolset import tool, Toolset
 from .tool.utils import find_tool_by_name
 from .types import LlmRequestParams, GenerateTextResponse, StreamTextResponseSync, StreamTextResponseAsync
+from .types.tool import ToolFn, ToolDef, RawToolDef, ToolLike
 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
 
@@ -87,7 +87,9 @@ class LLM:
         ) -> list[ToolMessage]:
         results = []
         for tool_call_tuple in tool_call_tuples:
-            id, function_name, function_arguments = tool_call_tuple
+            function_name = tool_call_tuple.function_name
+            function_arguments = tool_call_tuple.function_arguments
+
             if (target_tool := find_tool_by_name(tools, function_name)) is None:
                 logger.warning(f"Tool \"{function_name}\" not found, skipping execution.")
                 continue
@@ -101,7 +103,7 @@ class LLM:
             except Exception as e:
                 error = f"{type(e).__name__}: {str(e)}"
             results.append(ToolMessage(
-                id=id,
+                id=tool_call_tuple.id,
                 name=function_name,
                 arguments=function_arguments,
                 result=result,
@@ -115,7 +117,9 @@ class LLM:
         ) -> list[ToolMessage]:
         results = []
         for tool_call_tuple in tool_call_tuples:
-            id, function_name, function_arguments = tool_call_tuple
+            function_name = tool_call_tuple.function_name
+            function_arguments = tool_call_tuple.function_arguments
+
             if (target_tool := find_tool_by_name(tools, function_name)) is None:
                 logger.warning(f"Tool \"{function_name}\" not found, skipping execution.")
                 continue
@@ -129,7 +133,7 @@ class LLM:
             except Exception as e:
                 error = f"{type(e).__name__}: {str(e)}"
             results.append(ToolMessage(
-                id=id,
+                id=tool_call_tuple.id,
                 name=function_name,
                 arguments=function_arguments,
                 result=result,
@@ -146,10 +150,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 +162,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)):
@@ -195,7 +195,7 @@ class LLM:
         return returned_stream, full_message_queue
 
     async def stream_text(self, params: LlmRequestParams) -> StreamTextResponseAsync:
-        async def stream(response: CustomStreamWrapper) -> AsyncGenerator[TextChunk | ReasoningChunk | AudioChunk | ImageChunk | ToolCallChunk]:
+        async def stream(response: CustomStreamWrapper) -> AsyncGenerator[MessageChunk]:
             nonlocal message_collector
             async for chunk in response:
                 chunk = cast(LiteLlmModelResponseStream, chunk)
@@ -233,8 +233,12 @@ __all__ = [
     "Timeout",
 
     "LLM",
+    "LlmProviders",
     "LlmRequestParams",
 
+    "tool",
+    "Toolset",
+
     "ToolFn",
     "ToolDef",
     "RawToolDef",
@@ -250,6 +254,7 @@ __all__ = [
 
     "MessageChunk",
     "TextChunk",
+    "UsageChunk",
     "ReasoningChunk",
     "AudioChunk",
     "ImageChunk",

liteai_sdk/param_parser.py

@@ -1,7 +1,8 @@
 from typing import Any
 from litellm.types.utils import LlmProviders
-from .tool import prepare_tools
+from .tool.prepare import prepare_tools
 from .types import LlmRequestParams, ToolMessage
+from .types.tool import ToolLike
 
 ParsedParams = dict[str, Any]
 
@@ -14,8 +15,22 @@ class ParamParser:
         self._base_url = base_url
         self._api_key = api_key
 
+    @staticmethod
+    def _extract_tool_params(params: LlmRequestParams) -> list[ToolLike] | None:
+        if params.tools is None and params.toolsets is None:
+            return None
+        tools = []
+        if params.tools:
+            tools = params.tools
+        if params.toolsets:
+            for toolset in params.toolsets:
+                tools.extend(toolset.get_tool_methods())
+        return tools
+
     def _parse(self, params: LlmRequestParams) -> ParsedParams:
-        tools = params.tools and prepare_tools(params.tools)
+        extracted_tool_likes = self._extract_tool_params(params)
+        tools = extracted_tool_likes and prepare_tools(extracted_tool_likes)
+
         transformed_messages = []
         for message in params.messages:
             if type(message) is ToolMessage and\
@@ -45,4 +60,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/tool/__init__.py

@@ -1,310 +0,0 @@
-"""
-source: https://github.com/mozilla-ai/any-llm/blob/main/src/any_llm/tools.py
-"""
-
-import dataclasses
-import enum
-import inspect
-import types as _types
-from collections.abc import Callable, Mapping, Sequence
-from datetime import date, datetime, time
-from typing import Annotated as _Annotated, Literal as _Literal, is_typeddict as _is_typeddict,\
-                   Any, Awaitable, get_args, get_origin, get_type_hints
-from pydantic import BaseModel as PydanticBaseModel
-
-ToolFn = Callable[..., Any] | Callable[..., Awaitable[Any]]
-
-"""
-RawToolDef example:
-{
-    "name": "get_current_weather",
-    "description": "Get the current weather in a given location",
-    "parameters": {
-        "type": "object",
-        "properties": {
-            "location": {
-                "type": "string",
-                "description": "The city and state, e.g. San Francisco, CA",
-            },
-            "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
-        },
-        "required": ["location"],
-    }
-}
-"""
-RawToolDef = dict[str, Any]
-
-@dataclasses.dataclass
-class ToolDef:
-    name: str
-    description: str
-    execute: ToolFn
-
-ToolLike = ToolDef | RawToolDef | ToolFn
-
-def _python_type_to_json_schema(python_type: Any) -> dict[str, Any]:
-    """Convert Python type annotation to a JSON Schema for a parameter.
-
-    Supported mappings (subset tailored for LLM tool schemas):
-    - Primitives: str/int/float/bool -> string/integer/number/boolean
-    - bytes -> string with contentEncoding base64
-    - datetime/date/time -> string with format date-time/date/time
-    - list[T] / Sequence[T] / set[T] / frozenset[T] -> array with items=schema(T)
-      - set/frozenset include uniqueItems=true
-      - list without type args defaults items to string
-    - dict[K,V] / Mapping[K,V] -> object with additionalProperties=schema(V)
-      - dict without type args defaults additionalProperties to string
-    - tuple[T1, T2, ...] -> array with prefixItems per element and min/maxItems
-    - tuple[T, ...] -> array with items=schema(T)
-    - Union[X, Y] and X | Y -> oneOf=[schema(X), schema(Y)] (without top-level type)
-    - Optional[T] (Union[T, None]) -> schema(T) (nullability not encoded)
-    - Literal[...]/Enum -> enum with appropriate type inference when uniform
-    - TypedDict -> object with properties/required per annotations
-    - dataclass/Pydantic BaseModel -> object with nested properties inferred from fields
-    """
-    origin = get_origin(python_type)
-    args = get_args(python_type)
-
-    if _Annotated is not None and origin is _Annotated and len(args) >= 1:
-        python_type = args[0]
-        origin = get_origin(python_type)
-        args = get_args(python_type)
-
-    if python_type is Any:
-        return {"type": "string"}
-
-    primitive_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
-    if python_type in primitive_map:
-        return {"type": primitive_map[python_type]}
-
-    if python_type is bytes:
-        return {"type": "string", "contentEncoding": "base64"}
-    if python_type is datetime:
-        return {"type": "string", "format": "date-time"}
-    if python_type is date:
-        return {"type": "string", "format": "date"}
-    if python_type is time:
-        return {"type": "string", "format": "time"}
-
-    if python_type is list:
-        return {"type": "array", "items": {"type": "string"}}
-    if python_type is dict:
-        return {"type": "object", "additionalProperties": {"type": "string"}}
-
-    if origin is _Literal:
-        literal_values = list(args)
-        schema_lit: dict[str, Any] = {"enum": literal_values}
-        if all(isinstance(v, bool) for v in literal_values):
-            schema_lit["type"] = "boolean"
-        elif all(isinstance(v, str) for v in literal_values):
-            schema_lit["type"] = "string"
-        elif all(isinstance(v, int) and not isinstance(v, bool) for v in literal_values):
-            schema_lit["type"] = "integer"
-        elif all(isinstance(v, int | float) and not isinstance(v, bool) for v in literal_values):
-            schema_lit["type"] = "number"
-        return schema_lit
-
-    if inspect.isclass(python_type) and issubclass(python_type, enum.Enum):
-        enum_values = [e.value for e in python_type]
-        value_types = {type(v) for v in enum_values}
-        schema: dict[str, Any] = {"enum": enum_values}
-        if value_types == {str}:
-            schema["type"] = "string"
-        elif value_types == {int}:
-            schema["type"] = "integer"
-        elif value_types <= {int, float}:
-            schema["type"] = "number"
-        elif value_types == {bool}:
-            schema["type"] = "boolean"
-        return schema
-
-    if _is_typeddict(python_type):
-        annotations: dict[str, Any] = getattr(python_type, "__annotations__", {}) or {}
-        required_keys = set(getattr(python_type, "__required_keys__", set()))
-        td_properties: dict[str, Any] = {}
-        td_required: list[str] = []
-        for field_name, field_type in annotations.items():
-            td_properties[field_name] = _python_type_to_json_schema(field_type)
-            if field_name in required_keys:
-                td_required.append(field_name)
-        schema_td: dict[str, Any] = {
-            "type": "object",
-            "properties": td_properties,
-        }
-        if td_required:
-            schema_td["required"] = td_required
-        return schema_td
-
-    if inspect.isclass(python_type) and dataclasses.is_dataclass(python_type):
-        type_hints = get_type_hints(python_type)
-        dc_properties: dict[str, Any] = {}
-        dc_required: list[str] = []
-        for field in dataclasses.fields(python_type):
-            field_type = type_hints.get(field.name, Any)
-            dc_properties[field.name] = _python_type_to_json_schema(field_type)
-            if (
-                field.default is dataclasses.MISSING
-                and getattr(field, "default_factory", dataclasses.MISSING) is dataclasses.MISSING
-            ):
-                dc_required.append(field.name)
-        schema_dc: dict[str, Any] = {"type": "object", "properties": dc_properties}
-        if dc_required:
-            schema_dc["required"] = dc_required
-        return schema_dc
-
-    if inspect.isclass(python_type) and issubclass(python_type, PydanticBaseModel):
-        model_type_hints = get_type_hints(python_type)
-        pd_properties: dict[str, Any] = {}
-        pd_required: list[str] = []
-        model_fields = getattr(python_type, "model_fields", {})
-        for name, field_info in model_fields.items():
-            pd_properties[name] = _python_type_to_json_schema(model_type_hints.get(name, Any))
-            is_required = getattr(field_info, "is_required", None)
-            if callable(is_required) and is_required():
-                pd_required.append(name)
-        schema_pd: dict[str, Any] = {"type": "object", "properties": pd_properties}
-        if pd_required:
-            schema_pd["required"] = pd_required
-        return schema_pd
-
-    if origin in (list, Sequence, set, frozenset):
-        item_type = args[0] if args else Any
-        item_schema = _python_type_to_json_schema(item_type)
-        schema_arr: dict[str, Any] = {"type": "array", "items": item_schema or {"type": "string"}}
-        if origin in (set, frozenset):
-            schema_arr["uniqueItems"] = True
-        return schema_arr
-    if origin is tuple:
-        if not args:
-            return {"type": "array", "items": {"type": "string"}}
-        if len(args) == 2 and args[1] is Ellipsis:
-            return {"type": "array", "items": _python_type_to_json_schema(args[0])}
-        prefix_items = [_python_type_to_json_schema(a) for a in args]
-        return {
-            "type": "array",
-            "prefixItems": prefix_items,
-            "minItems": len(prefix_items),
-            "maxItems": len(prefix_items),
-        }
-
-    if origin in (dict, Mapping):
-        value_type = args[1] if len(args) >= 2 else Any
-        value_schema = _python_type_to_json_schema(value_type)
-        return {"type": "object", "additionalProperties": value_schema or {"type": "string"}}
-
-    typing_union = getattr(__import__("typing"), "Union", None)
-    if origin in (typing_union, _types.UnionType):
-        non_none_args = [a for a in args if a is not type(None)]
-        if len(non_none_args) > 1:
-            schemas = [_python_type_to_json_schema(arg) for arg in non_none_args]
-            return {"oneOf": schemas}
-        if non_none_args:
-            return _python_type_to_json_schema(non_none_args[0])
-        return {"type": "string"}
-
-    return {"type": "string"}
-
-def _parse_callable_properties(func: ToolFn) -> tuple[dict[str, dict[str, Any]], list[str]]:
-    sig = inspect.signature(func)
-    type_hints = get_type_hints(func)
-
-    properties: dict[str, dict[str, Any]] = {}
-    required: list[str] = []
-
-    for param_name, param in sig.parameters.items():
-        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
-            continue
-
-        annotated_type = type_hints.get(param_name, str)
-        param_schema = _python_type_to_json_schema(annotated_type)
-
-        type_name = getattr(annotated_type, "__name__", str(annotated_type))
-        properties[param_name] = {
-            **param_schema,
-            "description": f"Parameter {param_name} of type {type_name}",
-        }
-
-        if param.default == inspect.Parameter.empty:
-            required.append(param_name)
-
-    return properties, required
-
-def generate_tool_definition_from_callable(func: ToolFn) -> dict[str, Any]:
-    """Convert a Python callable to OpenAI tools format.
-
-    Args:
-        func: A Python callable (function) to convert to a tool
-
-    Returns:
-        Dictionary in OpenAI tools format
-
-    Raises:
-        ValueError: If the function doesn't have proper docstring or type annotations
-
-    Example:
-        >>> def get_weather(location: str, unit: str = "celsius") -> str:
-        ...     '''Get weather information for a location.'''
-        ...     return f"Weather in {location} is sunny, 25°{unit[0].upper()}"
-        >>>
-        >>> tool = generate_tool_definition_from_callable(get_weather)
-        >>> # Returns OpenAI tools format dict
-
-    """
-    if not func.__doc__:
-        msg = f"Function {func.__name__} must have a docstring"
-        raise ValueError(msg)
-
-    properties, required = _parse_callable_properties(func)
-    return {
-        "type": "function",
-        "function": {
-            "name": func.__name__,
-            "description": func.__doc__.strip(),
-            "parameters": {"type": "object", "properties": properties, "required": required},
-        },
-    }
-
-def generate_tool_definition_from_tool_def(tool_def: ToolDef) -> dict[str, Any]:
-    """Convert a ToolDef to OpenAI tools format.
-
-    Args:
-        tool_def: A ToolDef to convert to a tool
-
-    Returns:
-        Dictionary in OpenAI tools format
-
-    Example:
-        >>> tool_def = ToolDef(
-        ...     name="get_weather",
-        ...     description="Get weather information for a location.",
-        ...     execute=SomeFunction(),
-        ... )
-        >>> tool = generate_tool_definition_from_tool_def(tool_def)
-        >>> # Returns OpenAI tools format dict
-    """
-    properties, required = _parse_callable_properties(tool_def.execute)
-    return {
-        "type": "function",
-        "function": {
-            "name": tool_def.name,
-            "description": tool_def.description,
-            "parameters": {"type": "object", "properties": properties, "required": required},
-        },
-    }
-
-def generate_tool_definition_from_raw_tool_def(raw_tool_def: RawToolDef) -> dict[str, Any]:
-    return {
-        "type": "function",
-        "function": raw_tool_def,
-    }
-
-def prepare_tools(tools: list[ToolLike]) -> list[dict]:
-    tool_defs = []
-    for tool in tools:
-        if callable(tool):
-            tool_defs.append(generate_tool_definition_from_callable(tool))
-        elif isinstance(tool, ToolDef):
-            tool_defs.append(generate_tool_definition_from_tool_def(tool))
-        else:
-            tool_defs.append(generate_tool_definition_from_raw_tool_def(tool))
-    return tool_defs

liteai_sdk/tool/execute.py

@@ -3,7 +3,7 @@ import json
 from functools import singledispatch
 from typing import Any, Awaitable, Callable, cast
 from types import FunctionType, MethodType, CoroutineType
-from . import ToolDef
+from ..types.tool import ToolDef
 
 async def _coroutine_wrapper(awaitable: Awaitable[Any]) -> CoroutineType:
     return await awaitable

liteai_sdk/tool/prepare.py

@@ -0,0 +1,281 @@
+"""
+source: https://github.com/mozilla-ai/any-llm/blob/main/src/any_llm/tools.py
+"""
+
+import dataclasses
+import enum
+import inspect
+import types as _types
+from collections.abc import Mapping, Sequence
+from datetime import date, datetime, time
+from typing import Annotated as _Annotated, Literal as _Literal, is_typeddict as _is_typeddict,\
+                   Any, get_args, get_origin, get_type_hints
+from pydantic import BaseModel as PydanticBaseModel
+from ..types.tool import ToolFn, ToolDef, RawToolDef, ToolLike
+
+def _python_type_to_json_schema(python_type: Any) -> dict[str, Any]:
+    """Convert Python type annotation to a JSON Schema for a parameter.
+
+    Supported mappings (subset tailored for LLM tool schemas):
+    - Primitives: str/int/float/bool -> string/integer/number/boolean
+    - bytes -> string with contentEncoding base64
+    - datetime/date/time -> string with format date-time/date/time
+    - list[T] / Sequence[T] / set[T] / frozenset[T] -> array with items=schema(T)
+      - set/frozenset include uniqueItems=true
+      - list without type args defaults items to string
+    - dict[K,V] / Mapping[K,V] -> object with additionalProperties=schema(V)
+      - dict without type args defaults additionalProperties to string
+    - tuple[T1, T2, ...] -> array with prefixItems per element and min/maxItems
+    - tuple[T, ...] -> array with items=schema(T)
+    - Union[X, Y] and X | Y -> oneOf=[schema(X), schema(Y)] (without top-level type)
+    - Optional[T] (Union[T, None]) -> schema(T) (nullability not encoded)
+    - Literal[...]/Enum -> enum with appropriate type inference when uniform
+    - TypedDict -> object with properties/required per annotations
+    - dataclass/Pydantic BaseModel -> object with nested properties inferred from fields
+    """
+    origin = get_origin(python_type)
+    args = get_args(python_type)
+
+    if _Annotated is not None and origin is _Annotated and len(args) >= 1:
+        python_type = args[0]
+        origin = get_origin(python_type)
+        args = get_args(python_type)
+
+    if python_type is Any:
+        return {"type": "string"}
+
+    primitive_map = {str: "string", int: "integer", float: "number", bool: "boolean"}
+    if python_type in primitive_map:
+        return {"type": primitive_map[python_type]}
+
+    if python_type is bytes:
+        return {"type": "string", "contentEncoding": "base64"}
+    if python_type is datetime:
+        return {"type": "string", "format": "date-time"}
+    if python_type is date:
+        return {"type": "string", "format": "date"}
+    if python_type is time:
+        return {"type": "string", "format": "time"}
+
+    if python_type is list:
+        return {"type": "array", "items": {"type": "string"}}
+    if python_type is dict:
+        return {"type": "object", "additionalProperties": {"type": "string"}}
+
+    if origin is _Literal:
+        literal_values = list(args)
+        schema_lit: dict[str, Any] = {"enum": literal_values}
+        if all(isinstance(v, bool) for v in literal_values):
+            schema_lit["type"] = "boolean"
+        elif all(isinstance(v, str) for v in literal_values):
+            schema_lit["type"] = "string"
+        elif all(isinstance(v, int) and not isinstance(v, bool) for v in literal_values):
+            schema_lit["type"] = "integer"
+        elif all(isinstance(v, int | float) and not isinstance(v, bool) for v in literal_values):
+            schema_lit["type"] = "number"
+        return schema_lit
+
+    if inspect.isclass(python_type) and issubclass(python_type, enum.Enum):
+        enum_values = [e.value for e in python_type]
+        value_types = {type(v) for v in enum_values}
+        schema: dict[str, Any] = {"enum": enum_values}
+        if value_types == {str}:
+            schema["type"] = "string"
+        elif value_types == {int}:
+            schema["type"] = "integer"
+        elif value_types <= {int, float}:
+            schema["type"] = "number"
+        elif value_types == {bool}:
+            schema["type"] = "boolean"
+        return schema
+
+    if _is_typeddict(python_type):
+        annotations: dict[str, Any] = getattr(python_type, "__annotations__", {}) or {}
+        required_keys = set(getattr(python_type, "__required_keys__", set()))
+        td_properties: dict[str, Any] = {}
+        td_required: list[str] = []
+        for field_name, field_type in annotations.items():
+            td_properties[field_name] = _python_type_to_json_schema(field_type)
+            if field_name in required_keys:
+                td_required.append(field_name)
+        schema_td: dict[str, Any] = {
+            "type": "object",
+            "properties": td_properties,
+        }
+        if td_required:
+            schema_td["required"] = td_required
+        return schema_td
+
+    if inspect.isclass(python_type) and dataclasses.is_dataclass(python_type):
+        type_hints = get_type_hints(python_type)
+        dc_properties: dict[str, Any] = {}
+        dc_required: list[str] = []
+        for field in dataclasses.fields(python_type):
+            field_type = type_hints.get(field.name, Any)
+            dc_properties[field.name] = _python_type_to_json_schema(field_type)
+            if (
+                field.default is dataclasses.MISSING
+                and getattr(field, "default_factory", dataclasses.MISSING) is dataclasses.MISSING
+            ):
+                dc_required.append(field.name)
+        schema_dc: dict[str, Any] = {"type": "object", "properties": dc_properties}
+        if dc_required:
+            schema_dc["required"] = dc_required
+        return schema_dc
+
+    if inspect.isclass(python_type) and issubclass(python_type, PydanticBaseModel):
+        model_type_hints = get_type_hints(python_type)
+        pd_properties: dict[str, Any] = {}
+        pd_required: list[str] = []
+        model_fields = getattr(python_type, "model_fields", {})
+        for name, field_info in model_fields.items():
+            pd_properties[name] = _python_type_to_json_schema(model_type_hints.get(name, Any))
+            is_required = getattr(field_info, "is_required", None)
+            if callable(is_required) and is_required():
+                pd_required.append(name)
+        schema_pd: dict[str, Any] = {"type": "object", "properties": pd_properties}
+        if pd_required:
+            schema_pd["required"] = pd_required
+        return schema_pd
+
+    if origin in (list, Sequence, set, frozenset):
+        item_type = args[0] if args else Any
+        item_schema = _python_type_to_json_schema(item_type)
+        schema_arr: dict[str, Any] = {"type": "array", "items": item_schema or {"type": "string"}}
+        if origin in (set, frozenset):
+            schema_arr["uniqueItems"] = True
+        return schema_arr
+    if origin is tuple:
+        if not args:
+            return {"type": "array", "items": {"type": "string"}}
+        if len(args) == 2 and args[1] is Ellipsis:
+            return {"type": "array", "items": _python_type_to_json_schema(args[0])}
+        prefix_items = [_python_type_to_json_schema(a) for a in args]
+        return {
+            "type": "array",
+            "prefixItems": prefix_items,
+            "minItems": len(prefix_items),
+            "maxItems": len(prefix_items),
+        }
+
+    if origin in (dict, Mapping):
+        value_type = args[1] if len(args) >= 2 else Any
+        value_schema = _python_type_to_json_schema(value_type)
+        return {"type": "object", "additionalProperties": value_schema or {"type": "string"}}
+
+    typing_union = getattr(__import__("typing"), "Union", None)
+    if origin in (typing_union, _types.UnionType):
+        non_none_args = [a for a in args if a is not type(None)]
+        if len(non_none_args) > 1:
+            schemas = [_python_type_to_json_schema(arg) for arg in non_none_args]
+            return {"oneOf": schemas}
+        if non_none_args:
+            return _python_type_to_json_schema(non_none_args[0])
+        return {"type": "string"}
+
+    return {"type": "string"}
+
+def _parse_callable_properties(func: ToolFn) -> tuple[dict[str, dict[str, Any]], list[str]]:
+    sig = inspect.signature(func)
+    type_hints = get_type_hints(func)
+
+    properties: dict[str, dict[str, Any]] = {}
+    required: list[str] = []
+
+    for param_name, param in sig.parameters.items():
+        if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
+            continue
+
+        annotated_type = type_hints.get(param_name, str)
+        param_schema = _python_type_to_json_schema(annotated_type)
+
+        type_name = getattr(annotated_type, "__name__", str(annotated_type))
+        properties[param_name] = {
+            **param_schema,
+            "description": f"Parameter {param_name} of type {type_name}",
+        }
+
+        if param.default == inspect.Parameter.empty:
+            required.append(param_name)
+
+    return properties, required
+
+def generate_tool_definition_from_callable(func: ToolFn) -> dict[str, Any]:
+    """Convert a Python callable to OpenAI tools format.
+
+    Args:
+        func: A Python callable (function) to convert to a tool
+
+    Returns:
+        Dictionary in OpenAI tools format
+
+    Raises:
+        ValueError: If the function doesn't have proper docstring or type annotations
+
+    Example:
+        >>> def get_weather(location: str, unit: str = "celsius") -> str:
+        ...     '''Get weather information for a location.'''
+        ...     return f"Weather in {location} is sunny, 25°{unit[0].upper()}"
+        >>>
+        >>> tool = generate_tool_definition_from_callable(get_weather)
+        >>> # Returns OpenAI tools format dict
+
+    """
+    if not func.__doc__:
+        msg = f"Function {func.__name__} must have a docstring"
+        raise ValueError(msg)
+
+    properties, required = _parse_callable_properties(func)
+    return {
+        "type": "function",
+        "function": {
+            "name": func.__name__,
+            "description": inspect.cleandoc(func.__doc__),
+            "parameters": {"type": "object", "properties": properties, "required": required},
+        },
+    }
+
+def generate_tool_definition_from_tool_def(tool_def: ToolDef) -> dict[str, Any]:
+    """Convert a ToolDef to OpenAI tools format.
+
+    Args:
+        tool_def: A ToolDef to convert to a tool
+
+    Returns:
+        Dictionary in OpenAI tools format
+
+    Example:
+        >>> tool_def = ToolDef(
+        ...     name="get_weather",
+        ...     description="Get weather information for a location.",
+        ...     execute=SomeFunction(),
+        ... )
+        >>> tool = generate_tool_definition_from_tool_def(tool_def)
+        >>> # Returns OpenAI tools format dict
+    """
+    properties, required = _parse_callable_properties(tool_def.execute)
+    return {
+        "type": "function",
+        "function": {
+            "name": tool_def.name,
+            "description": tool_def.description,
+            "parameters": {"type": "object", "properties": properties, "required": required},
+        },
+    }
+
+def generate_tool_definition_from_raw_tool_def(raw_tool_def: RawToolDef) -> dict[str, Any]:
+    return {
+        "type": "function",
+        "function": raw_tool_def,
+    }
+
+def prepare_tools(tools: Sequence[ToolLike]) -> list[dict]:
+    tool_defs = []
+    for tool in tools:
+        if callable(tool):
+            tool_defs.append(generate_tool_definition_from_callable(tool))
+        elif isinstance(tool, ToolDef):
+            tool_defs.append(generate_tool_definition_from_tool_def(tool))
+        else:
+            tool_defs.append(generate_tool_definition_from_raw_tool_def(tool))
+    return tool_defs

liteai_sdk/tool/toolset.py

@@ -0,0 +1,18 @@
+import inspect
+from typing import Any, Callable, TypeVar
+from ..types.tool import ToolFn
+
+F = TypeVar("F", bound=Callable[..., Any])
+TOOL_FLAG = "__is_tool__"
+
+def tool(func: F) -> F:
+    setattr(func, TOOL_FLAG, True)
+    return func
+
+class Toolset:
+    def get_tool_methods(self) -> list[ToolFn]:
+        return [
+            method
+            for _, method in inspect.getmembers(self, predicate=inspect.ismethod)
+            if getattr(method, TOOL_FLAG, False)
+        ]

liteai_sdk/tool/utils.py

@@ -1,4 +1,4 @@
-from . import ToolFn, ToolDef, ToolLike
+from ..types.tool import ToolFn, ToolDef, ToolLike
 
 def find_tool_by_name(tools: list[ToolLike], name: str) -> ToolLike | None:
     for tool in tools:

liteai_sdk/types/__init__.py

@@ -1,9 +1,10 @@
 import asyncio
 import dataclasses
 import queue
-from typing import Any, Generator, Literal
+from typing import Any, Literal
 from collections.abc import AsyncGenerator, Generator
-from ..tool import ToolLike
+from .tool import ToolLike
+from ..tool.toolset import Toolset
 from .message import ChatMessage, AssistantMessage, ToolMessage, MessageChunk
 
 @dataclasses.dataclass
@@ -11,6 +12,7 @@ class LlmRequestParams:
     model: str
     messages: list[ChatMessage]
     tools: list[ToolLike] | None = None
+    toolsets: list[Toolset] | None = None
     tool_choice: Literal["auto", "required", "none"] = "auto"
     execute_tools: bool = False
 

liteai_sdk/types/message.py

@@ -3,13 +3,18 @@ from __future__ import annotations
 import json
 import dataclasses
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast
 from pydantic import BaseModel, ConfigDict, PrivateAttr, field_validator
-from litellm.types.utils import Message as LiteLlmMessage,\
-                                ModelResponseStream as LiteLlmModelResponseStream,\
-                                ChatCompletionAudioResponse,\
-                                ChatCompletionMessageToolCall,\
-                                ChatCompletionDeltaToolCall
+from litellm.types.utils import (
+    Message as LiteLlmMessage,
+    ModelResponse as LiteLlmModelResponse,
+    ModelResponseStream as LiteLlmModelResponseStream,
+    Choices as LiteLlmModelResponseChoices,
+    ChatCompletionAudioResponse,
+    ChatCompletionMessageToolCall,
+    ChatCompletionDeltaToolCall,
+    Usage as LiteLlmUsage
+)
 from litellm.types.llms.openai import (
     AllMessageValues,
     OpenAIMessageContent,
@@ -21,14 +26,14 @@ from litellm.types.llms.openai import (
     ChatCompletionToolMessage,
     ChatCompletionSystemMessage,
 )
-from ..tool import ToolLike
+from ..types.tool import ToolLike
 from ..tool.utils import find_tool_by_name
 from ..logger import logger
 
 if TYPE_CHECKING:
     from . import LlmRequestParams
 
-class ChatMessage(BaseModel, ABC):    
+class ChatMessage(BaseModel, ABC):
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
         validate_assignment=True,
@@ -46,7 +51,7 @@ class UserMessage(ChatMessage):
 
 class ToolMessage(ChatMessage):
     """
-    The `tool_def` field is ref to the target tool of the tool call, and 
+    The `tool_def` field is ref to the target tool of the tool call, and
     it will only be None when the target tool is not found
     """
     id: str
@@ -88,19 +93,27 @@ class ToolMessage(ChatMessage):
             content=content,
             tool_call_id=self.id)
 
-ToolCallTuple = tuple[str, str, str]
+class ToolCallTuple(NamedTuple):
+    id: str
+    function_name: str
+    function_arguments: str
+
 class AssistantMessage(ChatMessage):
     content: str | None = None
     reasoning_content: str | None = None
     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 +131,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":
@@ -143,7 +157,7 @@ class AssistantMessage(ChatMessage):
                function_name is None or\
                function_arguments is None:
                 return None
-            results.append((id, function_name, function_arguments))
+            results.append(ToolCallTuple(id, function_name, function_arguments))
         return results
 
     def get_partial_tool_messages(self) -> list[ToolMessage] | None:
@@ -164,22 +178,20 @@ class AssistantMessage(ChatMessage):
 
         results = []
         for tool_call in parsed_tool_calls:
-            id, name, arguments = tool_call
-
             tool_message = ToolMessage(
-                id=id,
-                name=name,
-                arguments=arguments,
+                id=tool_call.id,
+                name=tool_call.function_name,
+                arguments=tool_call.function_arguments,
                 result=None,
                 error=None)
 
             if has_tool_def:
                 assert self._request_params_ref and self._request_params_ref.tools
-                target_tool = find_tool_by_name(self._request_params_ref.tools, name)
+                target_tool = find_tool_by_name(self._request_params_ref.tools, tool_call.function_name)
                 if target_tool:
                     tool_message = tool_message.with_tool_def(target_tool)
                 else:
-                    logger.warning(f"Tool {name} not found in request params, "
+                    logger.warning(f"Tool {tool_call.function_name} not found in request params, "
                                     "tool_def will not be attached to the tool message")
 
             results.append(tool_message)
@@ -196,6 +208,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
@@ -215,7 +233,7 @@ class ToolCallChunk:
     arguments: str
     index: int
 
-MessageChunk = TextChunk | ReasoningChunk | AudioChunk | ImageChunk | ToolCallChunk
+MessageChunk = TextChunk | UsageChunk | ReasoningChunk | AudioChunk | ImageChunk | ToolCallChunk
 
 def openai_chunk_normalizer(
         chunk: LiteLlmModelResponseStream
@@ -239,4 +257,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/types/tool.py

@@ -0,0 +1,33 @@
+import dataclasses
+from collections.abc import Callable
+from typing import Any, Awaitable
+
+ToolFn = Callable[..., Any] | Callable[..., Awaitable[Any]]
+
+"""
+RawToolDef example:
+{
+    "name": "get_current_weather",
+    "description": "Get the current weather in a given location",
+    "parameters": {
+        "type": "object",
+        "properties": {
+            "location": {
+                "type": "string",
+                "description": "The city and state, e.g. San Francisco, CA",
+            },
+            "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+        },
+        "required": ["location"],
+    }
+}
+"""
+RawToolDef = dict[str, Any]
+
+@dataclasses.dataclass
+class ToolDef:
+    name: str
+    description: str
+    execute: ToolFn
+
+ToolLike = ToolDef | RawToolDef | ToolFn

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

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

liteai_sdk-0.4.0.dist-info/RECORD

@@ -0,0 +1,18 @@
+liteai_sdk/__init__.py,sha256=nXEHbRSYjpGFJnzmOB_cpx8NL12Eumw_IzY8UjD6GAM,10527
+liteai_sdk/debug.py,sha256=T7qIy1BeeUGlF40l9JCMMVn8pvvMJAEQeG4adQbOydA,69
+liteai_sdk/logger.py,sha256=99vJAQRKcu4CuHgZYAJ2zDQtGea6Bn3vJJrS-mtza7c,677
+liteai_sdk/param_parser.py,sha256=ae_aaOfwBNhR3QW7xxmjOf4D5ssS_9VN0IeRH5aBUYQ,2249
+liteai_sdk/stream.py,sha256=T9MLmgPC8te6qvSkBOh7vkl-I4OGCKuW1kEN6RkiCe0,3176
+liteai_sdk/tool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+liteai_sdk/tool/execute.py,sha256=AzYqX-oG89LLHrxouZtf7M7HohFf_WRtzqczlmw5Nks,2473
+liteai_sdk/tool/prepare.py,sha256=8JJiLev3pALu36gaqXHX9CtfCFWVM6nBbVBIrqm24po,11499
+liteai_sdk/tool/toolset.py,sha256=bl7qrrlBFz7HHpt8ZZlTBqjDin5MIOTtIz2o3H8kgRI,476
+liteai_sdk/tool/utils.py,sha256=A_4Jx1BacRX1KmK3t_9rDXrmSXj6v4fzNtqLsN12S0I,420
+liteai_sdk/types/__init__.py,sha256=WHp1YUOdINvv4shYBNo3xmEr_6B7boXvAAaefldkHbs,1071
+liteai_sdk/types/exceptions.py,sha256=hIGu06htOJxfEBAHx7KTvLQr0Y8GYnBLFJFlr_IGpDs,602
+liteai_sdk/types/message.py,sha256=vAp1uMv-WXFvEBhxFyw6rZt3IS-pLdANEuEFJJRJ8aY,9520
+liteai_sdk/types/tool.py,sha256=XbqbANr8D-FHDConmKCULoX3KfXkQXhCiSTHvVmKOl0,797
+liteai_sdk-0.4.0.dist-info/licenses/LICENSE,sha256=cTeVgQVJJcRdm1boa2P1FBnOeXfA_egV6s4PouyrCxg,1064
+liteai_sdk-0.4.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+liteai_sdk-0.4.0.dist-info/METADATA,sha256=6660sLr1g5VFHUcf57zfHBkYekJCWuOm58qOgBKlNr8,3023
+liteai_sdk-0.4.0.dist-info/RECORD,,

liteai_sdk-0.3.21.dist-info/RECORD

@@ -1,15 +0,0 @@
-liteai_sdk/__init__.py,sha256=KTHDeLyGGtVA-H8nJhJMBr_rKFV1h5cD5voZ8oPXI00,10608
-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/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,,