API / server.py
Severian's picture
Create server.py
3dc447d verified
raw
history blame
9.34 kB
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
@app.route('/v1/completions', methods=['POST'])
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)
@app.route('/v1/chat/completions', methods=['POST'])
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)