Update server.py

server.py

@@ -1,250 +1,137 @@
 import os
-from flask import Flask, request, jsonify, Response, stream_with_context
+from flask import Flask, request, jsonify
 import requests
 import logging
 from dotenv import load_dotenv
 import json
+from typing import Dict, Any, Tuple, Union
 
 # Load environment variables
 load_dotenv()
 
 app = Flask(__name__)
 
-ANTHROPIC_API_URL = os.getenv('ANTHROPIC_API_URL', 'https://relay.stagwellmarketingcloud.io/anthropic/v1/messages')
+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",
+def transform_to_anthropic_request(
+    request_data: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Transform OpenAI-style request to Anthropic format."""
+    if 'messages' in request_data:
+        # Chat completion format
+        return {
+            "model": request_data.get('model', 'claude-3-5-sonnet-20240620'),
             "max_tokens": request_data.get('max_tokens', 1024),
-            "messages": [{"role": "user", "content": request_data.get('prompt', '')}]
+            "messages": request_data['messages']
         }
-        
-        # 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
-            }
+    else:
+        # Regular completion format
+        return {
+            "model": request_data.get('model', 'claude-3-5-sonnet-20240620'),
+            "max_tokens": request_data.get('max_tokens', 1024),
+            "messages": [{
+                "role": "user",
+                "content": request_data.get('prompt', '')
+            }]
         }
-        
-        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):
+def make_anthropic_request(
+    request_data: Dict[str, Any]
+) -> Tuple[Dict[str, Any], int]:
+    """Make request to Anthropic API with proper error handling."""
     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)}")
+        anthropic_request = transform_to_anthropic_request(request_data)
+        logger.info(
+            f"Sending request to Anthropic API: {json.dumps(anthropic_request)}"
+        )
         
-        # Make streaming request to Anthropic
         response = requests.post(
             ANTHROPIC_API_URL,
             headers=headers,
             json=anthropic_request,
-            stream=True
+            timeout=30
         )
         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'
-            }
-        )
+        return response.json(), 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
+        error_msg = f"Error communicating with Anthropic API: {str(e)}"
+        if hasattr(e, 'response') and e.response is not None:
+            error_msg += f"\nResponse: {e.response.text}"
+        logger.error(error_msg)
+        return {"error": error_msg}, e.response.status_code if e.response else 500
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}")
-        return jsonify({"error": "An unexpected error occurred"}), 500
+        error_msg = f"Unexpected error: {str(e)}"
+        logger.error(error_msg)
+        return {"error": error_msg}, 500
 
+# Routes remain mostly unchanged, but use new helper functions
 @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)
+    """Handle completion-style requests."""
+    logger.info(f"Received completion request: {json.dumps(request.json)}")
+    
+    anthropic_response, status_code = make_anthropic_request(request.json)
+    if status_code != 200:
+        return jsonify(anthropic_response), status_code
+        
+    return jsonify({
+        "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
+        }
+    }), 200
 
 @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'
-    }
+    """Handle chat completion-style requests."""
+    logger.info(f"Received chat completion request: {json.dumps(request.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, status_code = make_anthropic_request(request.json)
+    if status_code != 200:
+        return jsonify(anthropic_response), status_code
         
-        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({
+        "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
+    }), 200
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=4224)