Update server.py

server.py

@@ -1,39 +1,35 @@
 import os
-from flask import Flask, request, jsonify
+import json
 import requests
 import logging
+from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
+from urllib.parse import parse_qs, urlparse
+from typing import Dict, Any, Tuple
 from dotenv import load_dotenv
-import json
-from typing import Dict, Any, Tuple, Union
 
 # Load environment variables
 load_dotenv()
 
-app = Flask(__name__)
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
+# Constants
 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 transform_to_anthropic_request(
-    request_data: Dict[str, Any]
-) -> Dict[str, Any]:
+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": request_data['messages']
         }
     else:
-        # Regular completion format
         return {
             "model": request_data.get('model', 'claude-3-5-sonnet-20240620'),
             "max_tokens": request_data.get('max_tokens', 1024),
@@ -79,59 +75,102 @@ def make_anthropic_request(
         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():
-    """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
+class AnthropicProxyHandler(SimpleHTTPRequestHandler):
+    def send_json_response(self, data: Dict[str, Any], status: int = 200):
+        """Helper method to send JSON responses."""
+        self.send_response(status)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(data).encode('utf-8'))
 
-@app.route('/v1/chat/completions', methods=['POST'])
-def chat_completions():
-    """Handle chat completion-style requests."""
-    logger.info(f"Received chat 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
+    def do_GET(self):
+        """Handle GET requests."""
+        if self.path == "/v1":
+            self.send_json_response({"status": True})
+        elif self.path == "/v1/models":
+            self.send_json_response({
+                "data": [{
+                    "id": "claude-3-5-sonnet-20240620",
+                    "object": "model",
+                    "created": 1706745202,
+                    "owned_by": "anthropic"
+                }]
+            })
+        else:
+            super().do_GET()
+
+    def do_POST(self):
+        """Handle POST requests."""
+        content_length = int(self.headers.get('Content-Length', 0))
+        request_body = self.rfile.read(content_length)
+        request_data = json.loads(request_body)
         
-    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
-        }
-    }), 200
+        if self.path == "/v1/completions":
+            logger.info(f"Received completion request: {json.dumps(request_data)}")
+            
+            response, status_code = make_anthropic_request(request_data)
+            if status_code != 200:
+                self.send_json_response(response, status_code)
+                return
+
+            completion_response = {
+                "id": "cmpl-" + response.get('id', 'default'),
+                "object": "text_completion",
+                "created": response.get('created', 0),
+                "choices": [{
+                    "text": response['content'][0]['text'],
+                    "index": 0,
+                    "finish_reason": "stop"
+                }],
+                "usage": {
+                    "prompt_tokens": -1,
+                    "completion_tokens": -1,
+                    "total_tokens": -1
+                }
+            }
+            self.send_json_response(completion_response)
+
+        elif self.path == "/v1/chat/completions":
+            logger.info(
+                f"Received chat completion request: {json.dumps(request_data)}"
+            )
+            
+            response, status_code = make_anthropic_request(request_data)
+            if status_code != 200:
+                self.send_json_response(response, status_code)
+                return
+
+            chat_response = {
+                "id": "chatcmpl-" + response.get('id', 'default'),
+                "object": "chat.completion",
+                "created": response.get('created', 0),
+                "choices": [{
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": response['content'][0]['text']
+                    },
+                    "finish_reason": "stop"
+                }],
+                "usage": {
+                    "prompt_tokens": -1,
+                    "completion_tokens": -1,
+                    "total_tokens": -1
+                }
+            }
+            self.send_json_response(chat_response)
+
+        else:
+            self.send_json_response(
+                {"error": "Invalid endpoint"}, 
+                status=404
+            )
+
+def run_server(port: int = 7860):
+    """Start the server on the specified port."""
+    server = ThreadingHTTPServer(("", port), AnthropicProxyHandler)
+    logger.info(f"Starting server on port {port}")
+    server.serve_forever()
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=4224)
+if __name__ == "__main__":
+    run_server()