# Formula Chat Client - OpenAI chat with official tools
# Uses MOONSHOT_BASE_URL and MOONSHOT_API_KEY for OpenAI client
import os
import json
import asyncio
import argparse
import httpx
from openai import AsyncOpenAI
class FormulaChatClient:
def __init__(self, moonshot_base_url: str, api_key: str):
self.openai = AsyncOpenAI(base_url=moonshot_base_url, api_key=api_key)
self.httpx = httpx.AsyncClient(
base_url=moonshot_base_url,
headers={"Authorization": f"Bearer {api_key}"},
timeout=30.0,
)
self.model = "kimi-k2.6"
async def get_tools(self, formula_uri: str):
response = await self.httpx.get(f"/formulas/{formula_uri}/tools")
return response.json().get("tools", [])
async def call_tool(self, formula_uri: str, function: str, args: dict):
response = await self.httpx.post(
f"/formulas/{formula_uri}/fibers",
json={"name": function, "arguments": json.dumps(args)},
)
fiber = response.json()
if fiber.get("status", "") == "succeeded":
return fiber["context"].get("output") or fiber["context"].get(
"encrypted_output"
)
if "error" in fiber:
return f"Error: {fiber['error']}"
if "error" in fiber.get("context", {}):
return f"Error: {fiber['context']['error']}"
if "output" in fiber.get("context", {}):
return f"Error: {fiber['context']['output']}"
return "Error: Unknown error"
async def handle_response(self, response, messages, all_tools, tool_to_uri):
message = response.choices[0].message
messages.append(message)
if not message.tool_calls:
print(f"\nAI Response: {message.content}")
return
print(f"\nAI decided to use {len(message.tool_calls)} tool(s):")
for call in message.tool_calls:
func_name = call.function.name
args = json.loads(call.function.arguments)
print(f"\nCalling tool: {func_name}")
print(f"Arguments: {json.dumps(args, ensure_ascii=False, indent=2)}")
uri = tool_to_uri.get(func_name)
if not uri:
raise ValueError(f"No URI found for tool {func_name}")
result = await self.call_tool(uri, func_name, args)
if len(result) > 100:
print(f"Tool result: {result[:100]}...") # limit the output length
else:
print(f"Tool result: {result}")
messages.append(
{"role": "tool", "tool_call_id": call.id, "content": result}
)
next_response = await self.openai.chat.completions.create(
model=self.model, messages=messages, tools=all_tools
)
await self.handle_response(next_response, messages, all_tools, tool_to_uri)
async def chat(self, question, messages, all_tools, tool_to_uri):
messages.append({"role": "user", "content": question})
response = await self.openai.chat.completions.create(
model=self.model, messages=messages, tools=all_tools
)
await self.handle_response(response, messages, all_tools, tool_to_uri)
async def close(self):
await self.httpx.aclose()
def normalize_formula_uri(uri: str) -> str:
"""Normalize formula URI with default namespace and tag"""
if "/" not in uri:
uri = f"moonshot/{uri}"
if ":" not in uri:
uri = f"{uri}:latest"
return uri
async def main():
parser = argparse.ArgumentParser(description="Chat with formula tools")
parser.add_argument(
"--formula",
action="append",
default=["moonshot/web-search:latest"],
help="Formula URIs",
)
parser.add_argument("--question", help="Question to ask")
args = parser.parse_args()
# Process and deduplicate formula URIs
raw_formulas = args.formula or ["moonshot/web-search:latest"]
normalized_formulas = [normalize_formula_uri(uri) for uri in raw_formulas]
unique_formulas = list(
dict.fromkeys(normalized_formulas)
) # Preserve order while deduping
print(f"Initialized formulas: {unique_formulas}")
moonshot_base_url = os.getenv("MOONSHOT_BASE_URL", "https://api.moonshot.cn/v1")
api_key = os.getenv("MOONSHOT_API_KEY")
if not api_key:
print("MOONSHOT_API_KEY required")
return
client = FormulaChatClient(moonshot_base_url, api_key)
# Load and validate tools
print("\nLoading tools from all formulas...")
all_tools = []
function_names = set()
tool_to_uri = {} # inverted index to the tool name
for uri in unique_formulas:
tools = await client.get_tools(uri)
print(f"\nTools from {uri}:")
for tool in tools:
func = tool.get("function", None)
if not func:
print(f"Skipping tool using type: {tool.get('type', 'unknown')}")
continue
func_name = func.get("name")
assert func_name, f"Tool missing name: {tool}"
assert (
func_name not in tool_to_uri
), f"ERROR: Tool '{func_name}' conflicts between {tool_to_uri.get(func_name)} and {uri}"
if func_name in function_names:
print(
f"ERROR: Duplicate function name '{func_name}' found across formulas"
)
print(f"Function {func_name} already exists in another formula")
await client.close()
return
function_names.add(func_name)
all_tools.append(tool)
tool_to_uri[func_name] = uri
print(f" - {func_name}: {func.get('description', 'N/A')}")
print(f"\nTotal unique tools loaded: {len(all_tools)}")
if not all_tools:
print("Warning: No tools found in any formula")
return
try:
messages = [
{
"role": "system",
"content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。",
}
]
if args.question:
print(f"\nUser: {args.question}")
await client.chat(args.question, messages, all_tools, tool_to_uri)
else:
print("Chat mode (type 'q' to quit)")
while True:
question = input("\nQ: ").strip()
if question.lower() == "q":
break
if question:
await client.chat(question, messages, all_tools, tool_to_uri)
finally:
await client.close()
if __name__ == "__main__":
asyncio.run(main())
