# 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.5"
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.ai/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": "You are Kimi, an AI assistant provided by Moonshot AI. You are proficient in Chinese and English conversations. You provide users with safe, helpful, and accurate answers. You will reject any questions involving terrorism, racism, or explicit content. Moonshot AI is a proper noun and should not be translated.",
}
]
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())
