You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
cbmc/specgen/util/deepseek_wrapper.py

128 lines
4.4 KiB

import json
import time
import os
import sys
import signal
import requests
def create_deepseek_config(messages):
config = {
"model": "deepseek-chat",
"temperature": 0.4,
"max_tokens": 4096,
"messages": []
}
config["messages"] = messages
return config
def handler(signum, frame):
raise Exception("end of time")
def request_deepseek_engine(config):
ret = None
max_retries = 5
retry_delay = 5
while ret is None and max_retries > 0:
try:
# Check if we're in the main thread before using signals
import threading
in_main_thread = threading.current_thread() is threading.main_thread()
# Only use signal.alarm on Unix-like systems and in main thread
if sys.platform != 'win32' and in_main_thread:
signal.signal(signal.SIGALRM, handler)
signal.alarm(100)
# Read API key from environment variable or file
api_key = os.getenv('DEEPSEEK_API_KEY')
if not api_key:
try:
with open("./api_key.txt", 'r') as f:
api_key = f.read().strip()
except FileNotFoundError:
raise Exception("DeepSeek API key not found. Set DEEPSEEK_API_KEY environment variable or create api_key.txt file")
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}"
}
response = requests.post(
"https://api.deepseek.com/v1/chat/completions",
headers=headers,
json=config,
timeout=100
)
if response.status_code == 200:
data = response.json()
# Format response to match OpenAI structure
ret = {
'choices': [{
'message': {
'role': data['choices'][0]['message']['role'],
'content': data['choices'][0]['message']['content']
}
}]
}
else:
print(f"API Error: {response.status_code} - {response.text}")
if response.status_code == 429: # Rate limit
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # Clear alarm before sleep
retry_delay = min(retry_delay * 2, 60) # Exponential backoff, max 60s
time.sleep(retry_delay)
max_retries -= 1
continue
else:
max_retries -= 1
if max_retries > 0:
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # Clear alarm before sleep
time.sleep(retry_delay)
continue
except requests.exceptions.ConnectionError as e:
print("Connection error. Waiting...")
print(e)
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # cancel alarm
time.sleep(retry_delay)
max_retries -= 1
retry_delay = min(retry_delay * 2, 60)
except requests.exceptions.Timeout as e:
print("Timeout error. Waiting...")
print(e)
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # cancel alarm
time.sleep(retry_delay)
max_retries -= 1
retry_delay = min(retry_delay * 2, 60)
except Exception as e:
print(e)
print("Unknown error. Waiting...")
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # cancel alarm
time.sleep(1)
max_retries -= 1
finally:
# Only clear alarm on platforms that support it and in main thread
if sys.platform != 'win32' and in_main_thread:
signal.alarm(0) # Always clear alarm
if ret is None:
# Return a fallback response if all retries failed
ret = {
'choices': [{
'message': {
'role': 'assistant',
'content': '```\n// Error: API request failed. Please check your API key and network connection.\n```'
}
}]
}
return ret