import os
from typing import *
import json
import httpx
from openai import OpenAI
client = OpenAI(
api_key=os.environ["MOONSHOT_API_KEY"], # 运行前请设置 MOONSHOT_API_KEY 环境变量
base_url="https://api.moonshot.cn/v1",
)
functions = [
{
"name": "search", # 函数的名称，请使用英文大小写字母、数据加上减号和下划线作为函数名称
"description": """
通过搜索引擎搜索互联网上的内容。
当你的知识无法回答用户提出的问题，或用户请求你进行联网搜索时，调用此工具。请从与用户的对话中提取用户想要搜索的内容作为 query 参数的值。
搜索结果包含网站的标题、网站的地址（URL）以及网站简介。
""", # 函数的介绍，在这里写上函数的具体作用以及使用场景，以便 Kimi 大模型能正确地选择使用哪些函数
"parameters": { # 使用 parameters 字段来定义函数接收的参数
"type": "object", # 固定使用 type: object 来使 Kimi 大模型生成一个 JSON Object 参数
"required": ["query"], # 使用 required 字段告诉 Kimi 大模型哪些参数是必填项
"properties": { # properties 中是具体的参数定义，你可以定义多个参数
"query": { # 在这里，key 是参数名称，value 是参数的具体定义
"type": "string", # 使用 type 定义参数类型
"description": """
用户搜索的内容，请从用户的提问或聊天上下文中提取。
""" # 使用 description 描述参数以便 Kimi 大模型更好地生成参数
}
}
}
}
]
def search_impl(query: str) -> List[Dict[str, Any]]:
"""
search_impl 使用搜索引擎对 query 进行搜索，目前主流的搜索引擎（例如 Bing）都提供了 API 调用方式，你可以自行选择
你喜欢的搜索引擎 API 进行调用，并将返回结果中的网站标题、网站链接、网站简介信息放置在一个 dict 中返回。
这里只是一个简单的示例，你可能需要编写一些鉴权、校验、解析的代码。
"""
r = httpx.get("https://your.search.api", params={"query": query})
return r.json()
def search(arguments: Dict[str, Any]) -> Any:
query = arguments["query"]
result = search_impl(query)
return {"result": result}
function_map = {
"search": search,
}
# ==========================================================================================================================================================
# tools 是 functions 的超集，因此我们可以通过已经定义的 functions 构造 tools，我们循环遍历每一个 function，并为其构造相应的 tool 格式；
# 同时，我们也生成相应的 tool_map。
# ==========================================================================================================================================================
tools = []
tool_map = {}
for function in functions:
tool = {
"type": "function",
"function": function,
}
tools.append(tool)
tool_map[function["name"]] = function_map[function["name"]]
messages = [
{"role": "system",
"content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"},
{"role": "user", "content": "请联网搜索 Context Caching，并告诉我它是什么。"} # 在提问中要求 Kimi 大模型联网搜索
]
finish_reason = None
# ==========================================================================================================================================================
# 在这里，我们将 finish_reason 值判断由 function_call 修改成 tool_calls
# ==========================================================================================================================================================
# while finish_reason is None or finish_reason == "function_call":
while finish_reason is None or finish_reason == "tool_calls":
completion = client.chat.completions.create(
model="kimi-k3",
messages=messages,
# ==========================================================================================================================================================
# 我们弃用 functions 参数，而是使用 tools 参数来启用工具调用
# ==========================================================================================================================================================
# function=functions,
tools=tools, # <-- 我们通过 tools 参数，将定义好的 tools 提交给 Kimi 大模型
)
choice = completion.choices[0]
finish_reason = choice.finish_reason
# ==========================================================================================================================================================
# 在这里，我们将原先 function_call 的执行逻辑替换成 tool_calls 的执行逻辑；
# 注意，由于 tool_calls 可能有多个，因此我们需要通过 for 循环逐个执行每个 tool_call。
# ==========================================================================================================================================================
# if finish_reason == "function_call":
# messages.append(choice.message)
# function_call_name = choice.message.function_call.name
# function_call_arguments = json.loads(choice.message.function_call.arguments)
# function_call = function_map[function_call_name]
# function_result = function_call(function_call_arguments)
# messages.append({
# "role": "function",
# "name": function_call_name,
# "content": json.dumps(function_result)
# })
if finish_reason == "tool_calls": # <-- 判断当前返回内容是否包含 tool_calls
messages.append(choice.message) # <-- 我们将 Kimi 大模型返回给我们的 assistant 消息也添加到上下文中，以便于下次请求时 Kimi 大模型能理解我们的诉求
for tool_call in choice.message.tool_calls: # <-- tool_calls 可能是多个，因此我们使用循环逐个执行
tool_call_name = tool_call.function.name
tool_call_arguments = json.loads(tool_call.function.arguments) # <-- arguments 是序列化后的 JSON Object，我们需要使用 json.loads 反序列化一下
tool_function = tool_map[tool_call_name] # <-- 通过 tool_map 快速找到需要执行哪个函数
tool_result = tool_function(tool_call_arguments)
# 使用函数执行结果构造一个 role=tool 的 message，以此来向模型展示工具调用的结果；
# 注意，我们需要在 message 中提供 tool_call_id 和 name 字段，以便 Kimi 大模型
# 能正确匹配到对应的 tool_call。
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"name": tool_call_name,
"content": json.dumps(tool_result), # <-- 我们约定使用字符串格式向 Kimi 大模型提交工具调用结果，因此在这里使用 json.dumps 将执行结果序列化成字符串
})
print(choice.message.content) # <-- 在这里，我们才将模型生成的回复返回给用户
