Create server.py

server.py

@@ -0,0 +1,250 @@
+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)