Source code for langchain_community.utils.ernie_functions

from typing import Literal, Optional, Type, TypedDict

from langchain_core.pydantic_v1 import BaseModel
from langchain_core.utils.json_schema import dereference_refs


class FunctionDescription(TypedDict):
    """代表Ernie API的可调用函数。"""

    name: str
    """函数的名称。"""
    description: str
    """一个函数的描述。"""
    parameters: dict
    """函数的参数。"""


class ToolDescription(TypedDict):
    """代表Ernie API的可调用函数。"""

    type: Literal["function"]
    function: FunctionDescription


def convert_pydantic_to_ernie_function(
    model: Type[BaseModel],
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
) -> FunctionDescription:
    """将Pydantic模型转换为Ernie API的函数描述。"""
    schema = dereference_refs(model.schema())
    schema.pop("definitions", None)
    return {
        "name": name or schema["title"],
        "description": description or schema["description"],
        "parameters": schema,
    }


def convert_pydantic_to_ernie_tool(
    model: Type[BaseModel],
    *,
    name: Optional[str] = None,
    description: Optional[str] = None,
) -> ToolDescription:
    """将Pydantic模型转换为Ernie API的函数描述。"""
    function = convert_pydantic_to_ernie_function(
        model, name=name, description=description
    )
    return {"type": "function", "function": function}