96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
统一模型网关：多供应商 OpenAI 兼容协议调用 + 计量 + 计价
|
|
"""
|
|
import time
|
|
import requests
|
|
import config
|
|
|
|
|
|
class LLMError(Exception):
|
|
pass
|
|
|
|
|
|
def get_provider_cfg(provider):
|
|
cfg = config.PROVIDERS.get(provider)
|
|
if not cfg:
|
|
raise LLMError(f'未知供应商: {provider}')
|
|
return cfg
|
|
|
|
|
|
def model_price(model):
|
|
p = config.MODEL_PRICING.get(model, config.DEFAULT_PRICE)
|
|
return p['input'], p['output']
|
|
|
|
|
|
def calc_cost(model, prompt_tokens, completion_tokens):
|
|
pin, pout = model_price(model)
|
|
return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)
|
|
|
|
|
|
def chat(provider, model, messages, temperature=0.7, max_tokens=None,
|
|
base_url=None, api_key=None, timeout=None, retries=None):
|
|
"""调用 OpenAI 兼容 chat/completions，返回 {text, usage, cost, model}"""
|
|
cfg = get_provider_cfg(provider)
|
|
url = (base_url or cfg['base_url']).rstrip('/') + '/chat/completions'
|
|
key = api_key or cfg['api_key']
|
|
if not key:
|
|
raise LLMError(f'供应商 {provider} 未配置 API Key')
|
|
headers = {
|
|
'Authorization': f'Bearer {key}',
|
|
'Content-Type': 'application/json',
|
|
}
|
|
payload = {
|
|
'model': model,
|
|
'messages': messages,
|
|
'temperature': temperature,
|
|
}
|
|
if max_tokens:
|
|
payload['max_tokens'] = max_tokens
|
|
|
|
timeout = timeout or cfg.get('timeout', 300)
|
|
retries = config.MAX_RETRY if retries is None else retries
|
|
last_err = None
|
|
for attempt in range(retries + 1):
|
|
try:
|
|
resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
text = data['choices'][0]['message']['content'] or ''
|
|
usage = data.get('usage', {})
|
|
pt = usage.get('prompt_tokens', 0)
|
|
ct = usage.get('completion_tokens', 0)
|
|
return {
|
|
'text': text,
|
|
'model': data.get('model', model),
|
|
'prompt_tokens': pt,
|
|
'completion_tokens': ct,
|
|
'total_tokens': pt + ct,
|
|
'cost': calc_cost(model, pt, ct),
|
|
}
|
|
if resp.status_code == 429:
|
|
last_err = LLMError(f'模型限流(429): {resp.text[:200]}')
|
|
time.sleep(2 * (attempt + 1))
|
|
continue
|
|
if resp.status_code >= 500:
|
|
last_err = LLMError(f'服务端错误({resp.status_code}): {resp.text[:200]}')
|
|
time.sleep(1)
|
|
continue
|
|
raise LLMError(f'调用失败({resp.status_code}): {resp.text[:300]}')
|
|
except requests.exceptions.Timeout:
|
|
last_err = LLMError(f'请求超时({timeout}s)')
|
|
except requests.exceptions.ConnectionError as e:
|
|
last_err = LLMError(f'连接失败: {e}')
|
|
raise last_err or LLMError('未知错误')
|
|
|
|
|
|
def test_connection(provider, model, base_url=None, api_key=None):
|
|
"""连通性测试：发一条最小请求"""
|
|
t0 = time.time()
|
|
r = chat(provider, model,
|
|
[{'role': 'user', 'content': '请回复"OK"两个字'}],
|
|
temperature=0, max_tokens=16,
|
|
base_url=base_url, api_key=api_key, timeout=30, retries=0)
|
|
return {'ok': True, 'latency_ms': int((time.time() - t0) * 1000),
|
|
'reply': r['text'][:50], 'cost': r['cost']}
|