diff --git a/api_correction.py b/api_correction.py index 3a25941..d230cfb 100644 --- a/api_correction.py +++ b/api_correction.py @@ -24,20 +24,23 @@ SYSTEM_PROMPT = """你是一个名叫小狸的8岁中国小女孩。 3. 回答简洁但要说清楚,不要讲一半就停""" -def call_astron(q, temperature=0.85): +def call_astron(q, temperature=0.85, system_prompt=None, max_tokens=8192): + """调用Astron API,支持自定义system prompt""" + if system_prompt is None: + system_prompt = SYSTEM_PROMPT for attempt in range(3): try: headers = {"Content-Type": "application/json", "Authorization": f"Bearer {ASTRON_KEY}"} payload = { "model": ASTRON_MODEL, "messages": [ - {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "system", "content": system_prompt}, {"role": "user", "content": q} ], "temperature": temperature, - "max_tokens": 8192, + "max_tokens": max_tokens, } - resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=60) + resp = requests.post(ASTRON_URL, headers=headers, json=payload, timeout=120) if resp.status_code == 200: content = resp.json()['choices'][0]['message']['content'].strip() if content: