import asyncio
import json
import os
import httpx
from openai import AsyncOpenAI
BASE_URL = "https://api.moonshot.cn/v1"
MODEL = "kimi-k3"
MAX_TOOL_ROUNDS = 8
SYSTEM_PROMPT = """你是行业研究助手。请先明确研究范围，再检索并交叉验证信息，最后输出简洁报告。
要求：
- 区分已确认事实、估算和推断；关键结论尽量由多个来源支持。
- 不得编造数据或来源；资料不足时说明搜索范围和缺口。
- 报告包含执行摘要、关键发现、风险与限制、来源列表。
- 使用与用户相同的语言回答。
"""
RESEARCH_PLAN_TOOL = {
"type": "function",
"function": {
"name": "build_research_plan",
"description": "根据行业主题、地区和时间范围生成研究计划",
"parameters": {
"type": "object",
"properties": {
"topic": {
"type": "string",
"description": "要研究的行业或主题",
},
"region": {
"type": "string",
"enum": ["中国", "全球", "美国", "欧洲"],
"description": "研究地区",
},
"start_year": {
"type": "integer",
"description": "研究起始年份",
},
"end_year": {
"type": "integer",
"description": "研究结束年份",
},
},
"required": ["topic", "region", "start_year", "end_year"],
"additionalProperties": False,
},
},
}
def build_research_plan(
topic: str, region: str, start_year: int, end_year: int
) -> str:
"""生成一个确定性的研究范围，作为自定义工具示例。"""
if start_year > end_year:
return json.dumps(
{"error": "start_year 不能大于 end_year"}, ensure_ascii=False
)
plan = {
"topic": topic,
"region": region,
"period": f"{start_year}-{end_year}",
"dimensions": ["市场规模与增速", "产业链与主要企业", "技术趋势", "政策与风险"],
"search_queries": [
f"{region} {topic} 市场规模 {start_year} {end_year}",
f"{region} {topic} 主要企业 技术趋势",
f"{region} {topic} 政策 风险",
],
}
return json.dumps(plan, ensure_ascii=False)
class IndustryResearchAgent:
def __init__(self) -> None:
api_key = os.environ["MOONSHOT_API_KEY"]
self.openai = AsyncOpenAI(api_key=api_key, base_url=BASE_URL)
self.http = httpx.AsyncClient(
base_url=BASE_URL,
headers={"Authorization": f"Bearer {api_key}"},
timeout=60.0,
)
async def load_formula(
self, formula_uri: str
) -> tuple[list[dict], dict[str, str]]:
response = await self.http.get(f"/formulas/{formula_uri}/tools")
response.raise_for_status()
tools = response.json().get("tools", [])
if not tools:
raise RuntimeError(f"Formula {formula_uri} 未返回任何工具")
tool_to_formula = {
tool["function"]["name"]: formula_uri
for tool in tools
if tool.get("type") == "function" and tool.get("function")
}
if not tool_to_formula:
raise RuntimeError(f"Formula {formula_uri} 未返回可调用的函数工具")
return tools, tool_to_formula
async def call_formula(
self, formula_uri: str, name: str, arguments: dict
) -> str:
response = await self.http.post(
f"/formulas/{formula_uri}/fibers",
json={"name": name, "arguments": json.dumps(arguments)},
)
response.raise_for_status()
fiber = response.json()
context = fiber.get("context", {})
if fiber.get("status") == "succeeded":
result = context.get("output") or context.get("encrypted_output") or ""
else:
result = fiber.get("error") or context.get("error") or "未知工具错误"
if isinstance(result, str):
return result
return json.dumps(result, ensure_ascii=False)
async def run(self, question: str) -> str:
official_tools, tool_to_formula = await self.load_formula(
"moonshot/web-search:latest"
)
tools = [RESEARCH_PLAN_TOOL, *official_tools]
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": question},
]
for _ in range(MAX_TOOL_ROUNDS):
response = await self.openai.chat.completions.create(
model=MODEL,
messages=messages,
tools=tools,
max_completion_tokens=8192,
)
choice = response.choices[0]
message = choice.message
# 追加完整 SDK message，保留 reasoning_content 和 tool_calls。
messages.append(message)
if choice.finish_reason == "length":
raise RuntimeError(
"模型输出被 max_completion_tokens 截断，请提高该参数或缩短工具结果"
)
if not message.tool_calls:
if choice.finish_reason != "stop":
raise RuntimeError(f"未预期的结束原因: {choice.finish_reason}")
if not message.content:
raise RuntimeError("模型未返回最终报告")
return message.content
for tool_call in message.tool_calls:
name = tool_call.function.name
try:
arguments = json.loads(tool_call.function.arguments or "{}")
if not isinstance(arguments, dict):
raise ValueError("工具参数必须是 JSON 对象")
if name == "build_research_plan":
result = build_research_plan(**arguments)
elif name in tool_to_formula:
result = await self.call_formula(
tool_to_formula[name], name, arguments
)
else:
raise ValueError(f"未知工具: {name}")
except Exception as exc:
result = json.dumps(
{"error": f"{type(exc).__name__}: {exc}"},
ensure_ascii=False,
)
# 即使单个工具失败，也返回对应结果并继续处理本轮其他调用。
messages.append(
{
"role": "tool",
"tool_call_id": tool_call.id,
"content": result,
}
)
raise RuntimeError(f"工具调用超过最大轮数 {MAX_TOOL_ROUNDS}")
async def close(self) -> None:
try:
await self.openai.close()
finally:
await self.http.aclose()
async def main() -> None:
agent = IndustryResearchAgent()
try:
report = await agent.run("调研 2024-2026 年中国人形机器人行业的发展情况")
print(report)
finally:
await agent.close()
if __name__ == "__main__":
asyncio.run(main())
