copilotkit__copilotkit/sdk-python/copilotkit/parameter.py

"""Parameter classes for CopilotKit"""

from typing import TypedDict, Optional, Literal, List, Union, cast, Any
from typing_extensions import NotRequired

class SimpleParameter(TypedDict):
    """Simple parameter class"""
    name: str
    description: NotRequired[str]
    required: NotRequired[bool]
    type: NotRequired[Literal[
        "number", 
        "boolean",
        "number[]", 
        "boolean[]"
    ]]

class ObjectParameter(TypedDict):
    """Object parameter class"""
    name: str
    description: NotRequired[str]
    required: NotRequired[bool]
    type: Literal["object", "object[]"]
    attributes: List['Parameter']

class StringParameter(TypedDict):
    """String parameter class"""
    name: str
    description: NotRequired[str]
    required: NotRequired[bool]
    type: Literal["string", "string[]"]
    enum: NotRequired[List[str]]

Parameter = Union[SimpleParameter, ObjectParameter, StringParameter]

def normalize_parameters(parameters: Optional[List[Parameter]]) -> List[Parameter]:
    """Normalize the parameters to ensure they have the correct type and format."""
    if parameters is None:
        return []
    return [_normalize_parameter(parameter) for parameter in parameters]

def _normalize_parameter(parameter: Parameter) -> Parameter:
    """Normalize a parameter to ensure it has the correct type and format."""
    if not "type" in parameter:
        cast(Any, parameter)['type'] = 'string'
    if not 'required' in parameter:
        parameter['required'] = True
    if not 'description' in parameter:
        parameter['description'] = ''

    if 'type' in parameter and (parameter['type'] == 'object' or parameter['type'] == 'object[]'):
        cast(Any, parameter)['attributes'] = normalize_parameters(parameter.get('attributes'))
    return parameter