Spaces:
Paused
Paused
import os | |
from flask import Flask, request, jsonify, Response, stream_with_context | |
import requests | |
import logging | |
from dotenv import load_dotenv | |
import json | |
# Load environment variables | |
load_dotenv() | |
app = Flask(__name__) | |
ANTHROPIC_API_URL = os.getenv('ANTHROPIC_API_URL', 'https://relay.stagwellmarketingcloud.io/anthropic/v1/messages') | |
API_KEY = os.getenv('ANTHROPIC_API_KEY') | |
# Configure logging | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
def handle_completion_request(request_data): | |
headers = { | |
'Authorization': f'Bearer {API_KEY}', | |
'Content-Type': 'application/json' | |
} | |
try: | |
# Transform to Anthropic format | |
anthropic_request = { | |
"model": "claude-3-sonnet-20240229", | |
"max_tokens": request_data.get('max_tokens', 1024), | |
"messages": [{"role": "user", "content": request_data.get('prompt', '')}] | |
} | |
# Make request to Anthropic | |
response = requests.post( | |
ANTHROPIC_API_URL, | |
headers=headers, | |
json=anthropic_request | |
) | |
response.raise_for_status() | |
# Get Anthropic response | |
anthropic_response = response.json() | |
# Transform to OpenAI format | |
openai_response = { | |
"id": "cmpl-" + anthropic_response.get('id', 'default'), | |
"object": "text_completion", | |
"created": anthropic_response.get('created', 0), | |
"choices": [{ | |
"text": anthropic_response['content'][0]['text'], | |
"index": 0, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": -1, | |
"completion_tokens": -1, | |
"total_tokens": -1 | |
} | |
} | |
return jsonify(openai_response), 200 | |
except requests.RequestException as e: | |
logger.error(f"Error communicating with Anthropic API: {str(e)}") | |
if e.response: | |
logger.error(f"Response content: {e.response.text}") | |
return jsonify({"error": "Error communicating with AI service"}), 500 | |
except Exception as e: | |
logger.error(f"Unexpected error: {str(e)}") | |
return jsonify({"error": "An unexpected error occurred"}), 500 | |
def stream_completion(request_data, is_chat=False): | |
headers = { | |
'Authorization': f'Bearer {API_KEY}', | |
'Content-Type': 'application/json', | |
'Accept': 'text/event-stream' | |
} | |
try: | |
# Transform to Anthropic format | |
anthropic_request = { | |
"model": "claude-3-sonnet-20240229", | |
"max_tokens": request_data.get('max_tokens', 1024), | |
"stream": True | |
} | |
if is_chat: | |
anthropic_request["messages"] = request_data.get('messages', []) | |
else: | |
anthropic_request["messages"] = [{"role": "user", "content": request_data.get('prompt', '')}] | |
logger.info(f"Sending streaming request to Anthropic: {json.dumps(anthropic_request)}") | |
# Make streaming request to Anthropic | |
response = requests.post( | |
ANTHROPIC_API_URL, | |
headers=headers, | |
json=anthropic_request, | |
stream=True | |
) | |
response.raise_for_status() | |
def generate(): | |
for line in response.iter_lines(): | |
if line: | |
line_text = line.decode('utf-8') | |
logger.info(f"Received line: {line_text}") | |
# Skip empty lines | |
if not line_text.strip(): | |
continue | |
# Handle SSE prefix | |
if line_text.startswith('data: '): | |
line_text = line_text[6:] # Remove 'data: ' prefix | |
try: | |
# Skip [DONE] message | |
if line_text.strip() == '[DONE]': | |
yield "data: [DONE]\n\n" | |
continue | |
data = json.loads(line_text) | |
logger.info(f"Parsed data: {json.dumps(data)}") | |
if is_chat: | |
chunk = { | |
"id": "chatcmpl-" + data.get('id', 'default'), | |
"object": "chat.completion.chunk", | |
"created": data.get('created', 0), | |
"model": "claude-3-sonnet-20240229", | |
"choices": [{ | |
"index": 0, | |
"delta": { | |
"role": "assistant", | |
"content": data.get('content', [{}])[0].get('text', '') | |
}, | |
"finish_reason": data.get('stop_reason') | |
}] | |
} | |
else: | |
chunk = { | |
"id": "cmpl-" + data.get('id', 'default'), | |
"object": "text_completion", | |
"created": data.get('created', 0), | |
"choices": [{ | |
"text": data.get('content', [{}])[0].get('text', ''), | |
"index": 0, | |
"finish_reason": data.get('stop_reason') | |
}] | |
} | |
yield f"data: {json.dumps(chunk)}\n\n" | |
if data.get('stop_reason'): | |
yield "data: [DONE]\n\n" | |
except json.JSONDecodeError as e: | |
logger.error(f"Error decoding JSON: {e}") | |
logger.error(f"Problematic line: {line_text}") | |
continue | |
except Exception as e: | |
logger.error(f"Error processing stream: {str(e)}") | |
continue | |
return Response( | |
stream_with_context(generate()), | |
mimetype='text/event-stream', | |
headers={ | |
'Cache-Control': 'no-cache', | |
'Transfer-Encoding': 'chunked' | |
} | |
) | |
except requests.RequestException as e: | |
logger.error(f"Error communicating with Anthropic API: {str(e)}") | |
if e.response: | |
logger.error(f"Response content: {e.response.text}") | |
return jsonify({"error": "Error communicating with AI service"}), 500 | |
except Exception as e: | |
logger.error(f"Unexpected error: {str(e)}") | |
return jsonify({"error": "An unexpected error occurred"}), 500 | |
def completions(): | |
request_data = request.json | |
if request_data.get('stream', False): | |
return stream_completion(request_data, is_chat=False) | |
return handle_completion_request(request_data) | |
def chat_completions(): | |
request_data = request.json | |
if request_data.get('stream', False): | |
return stream_completion(request_data, is_chat=True) | |
headers = { | |
'Authorization': f'Bearer {API_KEY}', | |
'Content-Type': 'application/json' | |
} | |
try: | |
anthropic_request = { | |
"model": "claude-3-sonnet-20240229", | |
"max_tokens": request_data.get('max_tokens', 1024), | |
"messages": request_data.get('messages', []) | |
} | |
response = requests.post( | |
ANTHROPIC_API_URL, | |
headers=headers, | |
json=anthropic_request | |
) | |
response.raise_for_status() | |
anthropic_response = response.json() | |
openai_response = { | |
"id": "chatcmpl-" + anthropic_response.get('id', 'default'), | |
"object": "chat.completion", | |
"created": anthropic_response.get('created', 0), | |
"choices": [{ | |
"index": 0, | |
"message": { | |
"role": "assistant", | |
"content": anthropic_response['content'][0]['text'] | |
}, | |
"finish_reason": "stop" | |
}], | |
"usage": { | |
"prompt_tokens": -1, | |
"completion_tokens": -1, | |
"total_tokens": -1 | |
} | |
} | |
return jsonify(openai_response), 200 | |
except requests.RequestException as e: | |
logger.error(f"Error communicating with Anthropic API: {str(e)}") | |
if e.response: | |
logger.error(f"Response content: {e.response.text}") | |
return jsonify({"error": "Error communicating with AI service"}), 500 | |
except Exception as e: | |
logger.error(f"Unexpected error: {str(e)}") | |
return jsonify({"error": "An unexpected error occurred"}), 500 | |
if __name__ == '__main__': | |
app.run(host='0.0.0.0', port=4224) | |