forked from bojieli/ai-agent-book
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathconfig.py
More file actions
223 lines (187 loc) · 8.04 KB
/
Copy pathconfig.py
File metadata and controls
223 lines (187 loc) · 8.04 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
"""
统一的 LLM 客户端配置。
默认使用 OpenAI(读取 OPENAI_API_KEY,模型 gpt-5.6-luna)。
也支持通过环境变量 LLM_PROVIDER 切换到 Moonshot / 火山方舟(ARK),
它们都兼容 OpenAI 的 Chat Completions + 工具调用接口。
export LLM_PROVIDER=openai # 默认
export LLM_PROVIDER=moonshot # 用 MOONSHOT_API_KEY
export LLM_PROVIDER=ark # 用 ARK_API_KEY,并需设置 ARK_MODEL
统一的 OpenRouter 兜底(fallback):
若所选 provider 自己的 Key 缺失,但设置了 OPENROUTER_API_KEY,则自动改走
OpenRouter(https://openrouter.ai/api/v1),并把模型名映射到 OpenRouter 命名:
gpt-* -> openai/gpt-*
claude-* -> anthropic/claude-opus-4.8
含 "/" -> 原样透传
其它 -> openai/gpt-5.6-luna
"""
import os
import time
from typing import Any
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()
OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
# 各提供商的默认配置:base_url / 环境变量名 / 默认模型
_PROVIDERS = {
"openai": {
"base_url": None, # 使用 SDK 默认
"key_env": "OPENAI_API_KEY",
"default_model": "gpt-5.6-luna",
},
"moonshot": {
"base_url": "https://api.moonshot.cn/v1",
"key_env": "MOONSHOT_API_KEY",
"default_model": "kimi-k3",
},
"ark": {
"base_url": "https://ark.cn-beijing.volces.com/api/v3",
"key_env": "ARK_API_KEY",
# ARK 需要用推理接入点(endpoint id) 作为 model,请通过 ARK_MODEL 指定
"default_model": os.getenv("ARK_MODEL", "doubao-seed-1-6-250615"),
},
"openrouter": {
"base_url": OPENROUTER_BASE_URL,
"key_env": "OPENROUTER_API_KEY",
"default_model": "openai/gpt-4o-mini",
},
}
API_TURNS = []
def _jsonable(value: Any) -> Any:
if hasattr(value, "model_dump"):
return _jsonable(value.model_dump(mode="json", exclude_none=True))
if isinstance(value, dict):
return {str(key): _jsonable(item) for key, item in value.items()}
if isinstance(value, (list, tuple)):
return [_jsonable(item) for item in value]
return value
def get_provider() -> str:
return os.getenv("LLM_PROVIDER", "openai").lower().strip()
def _to_openrouter_model(model: str) -> str:
"""把常见模型名映射到 OpenRouter 命名空间。"""
if not model:
return "openai/gpt-5.6-luna"
if "/" in model:
return model
if model.startswith("gpt-"):
return "openai/" + model
if model.startswith("claude-"):
return "anthropic/claude-opus-4.8"
return "openai/gpt-5.6-luna"
def _is_reasoning_model(model: str) -> bool:
"""gpt-5.x / o1·o3·o4 / kimi-k3 / *reasoner 等推理模型:不接受 temperature=0,
直连 gpt-5.x 还需组织实名且工具调用受限,故优先走 OpenRouter。"""
m = (model or "").lower()
return (m.startswith(("gpt-5", "o1", "o3", "o4"))
or m.startswith("kimi-k3")
or "reasoner" in m or "thinking" in m)
def _use_openrouter(cfg: dict) -> bool:
"""走 OpenRouter 的两种情形:
1) provider 自己的 Key 缺失、但有 OPENROUTER_API_KEY(统一兜底);
2) 目标是 gpt-5.x 且有 OPENROUTER_API_KEY —— 直连 gpt-5.x 需组织实名、
且 /chat/completions 工具调用受限,故即便有 OPENAI_API_KEY 也优先 OpenRouter。"""
if not os.getenv("OPENROUTER_API_KEY"):
return False
if not os.getenv(cfg["key_env"]):
return True
model = os.getenv("LLM_MODEL") or cfg["default_model"]
return (model or "").lower().startswith("gpt-5")
def get_model() -> str:
"""允许用 LLM_MODEL 覆盖默认模型;OpenRouter 兜底路径下映射模型名。"""
provider = get_provider()
if provider not in _PROVIDERS:
raise ValueError(f"未知的 LLM_PROVIDER: {provider}")
cfg = _PROVIDERS[provider]
model = os.getenv("LLM_MODEL") or cfg["default_model"]
if _use_openrouter(cfg):
return _to_openrouter_model(model)
return model
def get_client() -> OpenAI:
provider = get_provider()
if provider not in _PROVIDERS:
raise ValueError(f"未知的 LLM_PROVIDER: {provider}")
cfg = _PROVIDERS[provider]
if _use_openrouter(cfg):
return OpenAI(api_key=os.getenv("OPENROUTER_API_KEY"), base_url=OPENROUTER_BASE_URL)
api_key = os.getenv(cfg["key_env"])
if not api_key:
raise RuntimeError(
f"环境变量 {cfg['key_env']} 未设置,也未设置 OPENROUTER_API_KEY。"
f"请参考 env.example 配置其一(OpenRouter 可作为统一兜底)后重试。"
)
kwargs = {"api_key": api_key}
if cfg["base_url"]:
kwargs["base_url"] = cfg["base_url"]
return OpenAI(**kwargs)
def record_completion(client: OpenAI, *, kind: str, **request: Any):
"""Execute and retain a credential-free raw request/response receipt."""
started = time.time()
response = client.chat.completions.create(**request)
API_TURNS.append({
"kind": kind,
"provider": get_provider(),
"endpoint": get_backend_metadata()["endpoint"],
"request": _jsonable(request),
"response": response.model_dump(mode="json", exclude_none=True),
"elapsed_seconds": round(time.time() - started, 6),
})
return response
def reset_api_turns() -> None:
API_TURNS.clear()
def get_api_turns() -> list[dict]:
return list(API_TURNS)
def get_backend_metadata() -> dict[str, Any]:
provider = get_provider()
cfg = _PROVIDERS[provider]
if _use_openrouter(cfg):
base_url = OPENROUTER_BASE_URL
key_env = "OPENROUTER_API_KEY"
routed_provider = "openrouter"
else:
base_url = cfg["base_url"] or "https://api.openai.com/v1"
key_env = cfg["key_env"]
routed_provider = provider
return {
"configured_provider": provider,
"routed_provider": routed_provider,
"model": get_model(),
"endpoint": f"{base_url}/chat/completions",
"credential_source_env": key_env,
"credential_value_recorded": False,
}
def usage_summary() -> dict[str, Any]:
prompt = completion = total = 0
native_cost = 0.0
native_cost_count = 0
for turn in API_TURNS:
usage = turn.get("response", {}).get("usage") or {}
prompt += int(usage.get("prompt_tokens") or 0)
completion += int(usage.get("completion_tokens") or 0)
total += int(usage.get("total_tokens") or 0)
if usage.get("cost") is not None:
native_cost += float(usage["cost"])
native_cost_count += 1
return {
"prompt_tokens": prompt,
"completion_tokens": completion,
"total_tokens": total or prompt + completion,
"provider_reported_cost_usd": round(native_cost, 9) if native_cost_count else None,
"provider_reported_cost_observations": native_cost_count,
"cost_qualification": (
"provider-native usage.cost summed across calls"
if native_cost_count else "provider did not expose monetary cost; no price was guessed"
),
}
# 全部 LLM 调用统一使用低温度,保证结果可复现;
# 但推理模型(gpt-5.x / o 系列 / kimi-k3 等)只接受默认 temperature=1,
# 故按当前解析出的模型自动选择默认温度(可用 LLM_TEMPERATURE 显式覆盖)。
def _default_temperature() -> str:
provider = get_provider()
cfg = _PROVIDERS.get(provider, _PROVIDERS["openai"])
model = os.getenv("LLM_MODEL") or cfg["default_model"]
return "1" if _is_reasoning_model(model) else "0"
def get_temperature() -> float:
"""在调用时按当前解析出的模型选择温度,使 CLI/env 的 --model/--provider
覆盖生效。原来的模块级 TEMPERATURE 常量在 import 时就被固定,而 demo.py 在
import 之后才设置 LLM_MODEL/LLM_PROVIDER,导致温度停留在默认模型的值
(例如把非推理模型误用 temperature=1,破坏了本文件追求的可复现性)。"""
return float(os.getenv("LLM_TEMPERATURE", _default_temperature()))