import json
from typing import Optional, Any

from langchain.schema import AIMessage
from langchain.schema.runnable import RunnableSerializable, RunnableConfig
from pydantic import BaseModel


class FunctionCall(BaseModel):
    name: str
    arguments: dict[str, Any]


class ParseFunctionCall(RunnableSerializable[AIMessage, FunctionCall]):
    """
    Parse OpenAI function call from an AI Message. Used in conjunction with
    functions bound to the LLM request.
    """

    def invoke(
        self,
        input: AIMessage,
        config: Optional[RunnableConfig] = None,
        **kwargs: Any,
    ) -> FunctionCall:
        if isinstance(input, dict):
            input = AIMessage(**input)

        additional_kwargs = input.additional_kwargs
        if "function_call" in additional_kwargs:
            function_call = additional_kwargs["function_call"]
            if isinstance(function_call, str):
                function_call = json.loads(function_call)

            if isinstance(function_call["arguments"], str):
                function_call["arguments"] = json.loads(function_call["arguments"])

            return FunctionCall(**function_call)
        else:
            raise ValueError("No function call found in input.")

    async def ainvoke(
        self,
        input: AIMessage,
        config: Optional[RunnableConfig] = None,
        **kwargs: Any,
    ) -> FunctionCall:
        if isinstance(input, dict):
            input = AIMessage(**input)

        additional_kwargs = input.additional_kwargs
        if "function_call" in additional_kwargs:
            function_call = additional_kwargs["function_call"]
            if isinstance(function_call, str):
                function_call = json.loads(function_call)

            if isinstance(function_call["arguments"], str):
                function_call["arguments"] = json.loads(function_call["arguments"])

            return FunctionCall(**function_call)
        else:
            raise ValueError("No function call found in input.")

langchain.schema.AIMessage
[((682, 700), 'langchain.schema.AIMessage', 'AIMessage', ([], {}), '(**input)\n', (691, 700), False, 'from langchain.schema import AIMessage\n'), ((1460, 1478), 'langchain.schema.AIMessage', 'AIMessage', ([], {}), '(**input)\n', (1469, 1478), False, 'from langchain.schema import AIMessage\n'), ((945, 970), 'json.loads', 'json.loads', (['function_call'], {}), '(function_call)\n', (955, 970), False, 'import json\n'), ((1077, 1115), 'json.loads', 'json.loads', (["function_call['arguments']"], {}), "(function_call['arguments'])\n", (1087, 1115), False, 'import json\n'), ((1723, 1748), 'json.loads', 'json.loads', (['function_call'], {}), '(function_call)\n', (1733, 1748), False, 'import json\n'), ((1855, 1893), 'json.loads', 'json.loads', (["function_call['arguments']"], {}), "(function_call['arguments'])\n", (1865, 1893), False, 'import json\n')]