Repository: cpacker/MemGPT
File: letta/functions/schema_generator.py (Python)

Code before:
        assert optional_length(py_type) == 1, f"Optional type must have exactly one type argument, but got {py_type}"

        # Extract and map the inner type
        return type_to_json_schema_type(type_args[0])

    # Mapping of Python types to JSON schema types
    type_map = {
        # Basic types
        int: "integer",
        str: "string",
        bool: "boolean",
        float: "number",
        # Collections
        List[str]: "array",
        List[int]: "array",
        list: "array",
        tuple: "array",
        set: "array",
        # Dictionaries
        dict: "object",
        Dict[str, Any]: "object",
        # Special types
        None: "null",
        type(None): "null",
        # Optional types
        Optional[str]: "string",
        Union[str, None]: "string",
    }
    if py_type not in type_map:
        raise ValueError(f"Python type {py_type} has no corresponding JSON schema type - full map: {type_map}")

    return type_map.get(py_type, "string")  # Default to "string" if type not in map


def pydantic_model_to_open_ai(model):
    schema = model.model_json_schema()
    docstring = parse(model.__doc__ or "")
    parameters = {k: v for k, v in schema.items() if k not in ("title", "description")}
    for param in docstring.params:
        if (name := param.arg_name) in parameters["properties"] and (description := param.description):
            if "description" not in parameters["properties"][name]:
                parameters["properties"][name]["description"] = description

    parameters["required"] = sorted(k for k, v in parameters["properties"].items() if "default" not in v)

    if "description" not in schema:
        if docstring.short_description:
            schema["description"] = docstring.short_description
        else:
            raise


Review comment (suggestion):
@mattzh72 maybe it makes sense to kill this?
`Optional[str]: "string",`

because of the code above? `if is_optional(py_type):`

Code after:
        assert optional_length(py_type) == 1, f"Optional type must have exactly one type argument, but got {py_type}"

        # Extract and map the inner type
        return type_to_json_schema_type(type_args[0])

    # Mapping of Python types to JSON schema types
    type_map = {
        # Basic types
        int: "integer",
        str: "string",
        bool: "boolean",
        float: "number",
        # Collections
        List[str]: "array",
        List[int]: "array",
        list: "array",
        tuple: "array",
        set: "array",
        # Dictionaries
        dict: "object",
        Dict[str, Any]: "object",
        # Special types
        None: "null",
        type(None): "null",
        # Optional types
        # Optional[str]: "string",  # NOTE: caught above ^
        Union[str, None]: "string",
    }
    if py_type not in type_map:
        raise ValueError(f"Python type {py_type} has no corresponding JSON schema type - full map: {type_map}")

    return type_map.get(py_type, "string")  # Default to "string" if type not in map


def pydantic_model_to_open_ai(model):
    schema = model.model_json_schema()
    docstring = parse(model.__doc__ or "")
    parameters = {k: v for k, v in schema.items() if k not in ("title", "description")}
    for param in docstring.params:
        if (name := param.arg_name) in parameters["properties"] and (description := param.description):
            if "description" not in parameters["properties"][name]:
                parameters["properties"][name]["description"] = description

    parameters["required"] = sorted(k for k, v in parameters["properties"].items() if "default" not in v)

    if "description" not in schema:
        if docstring.short_description:
            schema["description"] = docstring.short_description
        else:
            raise