133 lines
5.0 KiB
Python
133 lines
5.0 KiB
Python
import json
|
|
import time
|
|
from typing import Dict, Any, Optional
|
|
|
|
import requests
|
|
|
|
|
|
class APIClient:
|
|
"""DeepSeek API 客户端，含重试逻辑。"""
|
|
|
|
def __init__(self, api_key: str, model: str = 'deepseek-v4-flash',
|
|
base_url: str = 'https://api.deepseek.com/chat/completions',
|
|
max_retries: int = 3, timeout: int = 120,
|
|
max_tokens: int = 32768):
|
|
self.api_key = api_key
|
|
self.model = model
|
|
self.base_url = base_url
|
|
self.max_retries = max_retries
|
|
self.timeout = timeout
|
|
self.max_tokens = max_tokens
|
|
|
|
def generate(self, prompt: str) -> Dict[str, Any]:
|
|
"""发送 prompt 并返回 AI 生成的结果。"""
|
|
system_prompt = (
|
|
"你是COBOL程序的测试数据生成专家。"
|
|
"请严格按照提供的规则，生成符合格式要求的测试数据。"
|
|
"输出必须是可被json.loads()直接解析的JSON，不要包裹在```json```代码块中。"
|
|
"不要在JSON前后添加任何说明文字。"
|
|
)
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": prompt},
|
|
]
|
|
|
|
last_error = None
|
|
|
|
for attempt in range(1, self.max_retries + 1):
|
|
try:
|
|
response = self._call_api(messages)
|
|
choice = response['choices'][0]
|
|
content = choice['message']['content']
|
|
|
|
if choice.get('finish_reason') == 'length':
|
|
# 输出被截断：追加"压缩输出"指示后重试，避免再次浪费全部 max_tokens
|
|
last_error = RuntimeError(
|
|
"API输出被截断（finish_reason=length），已追加压缩指示重试。"
|
|
)
|
|
error_msg = (
|
|
"前回の出力は長すぎて途中で切れました（finish_reason=length）。"
|
|
"レコード数・フィールド数を最小限にし、FILLERなどの冗長な項目は省略して、"
|
|
"必ず完結したJSONのみを出力してください。コードブロック（```）で囲まないでください。"
|
|
)
|
|
messages.append({"role": "assistant", "content": content})
|
|
messages.append({"role": "user", "content": error_msg})
|
|
continue
|
|
|
|
data = self._parse_json(content)
|
|
if data is not None:
|
|
return data
|
|
|
|
last_error = RuntimeError("API返回内容无法解析为JSON")
|
|
error_msg = (
|
|
f"前回の出力は有効なJSONではありませんでした。"
|
|
f"必ず有効なJSONのみを出力してください。"
|
|
f"コードブロック（```）で囲まないでください。"
|
|
)
|
|
messages.append({"role": "assistant", "content": content})
|
|
messages.append({"role": "user", "content": error_msg})
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
last_error = e
|
|
if attempt < self.max_retries:
|
|
time.sleep(2 ** attempt)
|
|
continue
|
|
|
|
raise RuntimeError(
|
|
f"API调用失败，已重试{self.max_retries}次。"
|
|
f"最后错误: {last_error}"
|
|
)
|
|
|
|
def _call_api(self, messages: list) -> dict:
|
|
"""单次 API 调用。"""
|
|
headers = {
|
|
'Authorization': f'Bearer {self.api_key}',
|
|
'Content-Type': 'application/json',
|
|
}
|
|
payload = {
|
|
'model': self.model,
|
|
# deepseek-v4 默认开启思考模式，会把 max_tokens 几乎全部消耗在
|
|
# reasoning_content 上，导致 content 为空/截断。测试数据生成是
|
|
# 结构化 JSON 输出，关闭思考模式后全部 token 用于正文。
|
|
'thinking': {'type': 'disabled'},
|
|
'messages': messages,
|
|
'temperature': 0.3,
|
|
'max_tokens': self.max_tokens,
|
|
}
|
|
|
|
resp = requests.post(
|
|
self.base_url,
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=self.timeout,
|
|
proxies={"http": None, "https": None},
|
|
)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
@staticmethod
|
|
def _parse_json(text: str) -> Optional[dict]:
|
|
"""尝试从文本中提取 JSON。"""
|
|
text = text.strip()
|
|
|
|
if text.startswith('```json'):
|
|
text = text[7:]
|
|
if text.startswith('```'):
|
|
text = text[3:]
|
|
if text.endswith('```'):
|
|
text = text[:-3]
|
|
text = text.strip()
|
|
|
|
try:
|
|
return json.loads(text)
|
|
except json.JSONDecodeError:
|
|
start = text.find('{')
|
|
end = text.rfind('}')
|
|
if start >= 0 and end > start:
|
|
try:
|
|
return json.loads(text[start:end + 1])
|
|
except json.JSONDecodeError:
|
|
pass
|
|
return None
|