from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool, render_text_description
from ChatGLM_new import tongyi_llm
from typing import Any, Dict, Optional, TypedDict
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables import RunnablePassthrough
@tool
def multiply(x: float, y: float) -> float:
    """Multiply two numbers together."""
    return x * y


@tool
def add(x: int, y: int) -> int:
    "Add two numbers."
    return x + y


tools = [multiply, add]

# Let's inspect the tools
for t in tools:
    print("--")
    print(t.name)
    print(t.description)
    print(t.args)
print("--------------------------")
rendered_tools = render_text_description(tools)
print(rendered_tools)


system_prompt = f"""\
您是有权访问以下一组工具的助手。
以下是每个工具的名称和说明：

{rendered_tools}

给定用户输入，返回要使用的工具的名称和输入。
将响应作为包含"name"和"arguments"键的 JSON blob 返回。

"参数"应该是字典，键对应
参数名称和与请求值对应的值。
"""
print("----------------------add-----------------------------")
prompt = ChatPromptTemplate.from_messages(
    [("system", system_prompt), ("user", "{input}")]
)
chain = prompt | tongyi_llm
message = chain.invoke({"input": "什么是 3 加 1132"})

# Let's take a look at the output from the model
# if the model is an LLM (not a chat model), the output will be a string.
if isinstance(message, str):
    print(message)
else:  # Otherwise it's a chat model
    print(message.content)
print("---------------------multiply------------------------------")
chain = prompt | tongyi_llm | JsonOutputParser()
print(chain.invoke({"input": "what's thirteen times 4"}))

print("---------------------invoke_tool------------------------------")

class ToolCallRequest(TypedDict):
    """显示 invoke_tool 函数的输入的类型化字典."""

    name: str
    arguments: Dict[str, Any]


def invoke_tool(
    tool_call_request: ToolCallRequest, config: Optional[RunnableConfig] = None
):
    """我们可以用来执行工具调用的函数.

    参数：
        tool_call_request：包含键名称和参数的字典。
            该名称必须与现有工具的名称匹配。
            参数是该工具的参数。
        config：这是LangChain使用的配置信息，包含
            诸如回调、元数据等。请参阅有关 RunnableConfig 的 LCEL 文档。

    返回：
        所请求工具的输出
    """
    tool_name_to_tool = {tool.name: tool for tool in tools}
    name = tool_call_request["name"]
    requested_tool = tool_name_to_tool[name]
    return requested_tool.invoke(tool_call_request["arguments"], config=config)


print(invoke_tool({"name": "multiply", "arguments": {"x": 3, "y": 5}}))
print("---------------------计算------------------------------")
chain = prompt | tongyi_llm | JsonOutputParser() | invoke_tool
#print(chain.invoke({"input": "what's thirteen times 4.14137281"}))

print("---------------------输出一个json结果------------------------------")
chain = (
    prompt | tongyi_llm | JsonOutputParser() | RunnablePassthrough.assign(output=invoke_tool)
)
print(chain.invoke({"input": "what's thirteen times 4.14137281"}))