Create app.py

app.py

@@ -0,0 +1,89 @@
+import os
+from flask import Flask, request, Response, jsonify
+import requests
+import json
+import random
+from helper import create_jwt, github_username_zed_userid_list
+
+app = Flask(__name__)
+
+@app.route('/hf/v1/chat/completions', methods=['POST'])
+def chat():
+    # Get the payload from the request
+    payload = request.json
+
+    # Get the model from the payload, defaulting to "claude-3-5-sonnet-20240620"
+    model = payload.get('model', 'claude-3-5-sonnet-20240620')
+
+    # Prepare the request for the LLM API
+    url = "https://llm.zed.dev/completion?"
+    
+    llm_payload = {
+        "provider": "anthropic",
+        "model": model,
+        "provider_request": {
+            "model": model,
+            "max_tokens": payload.get('max_tokens', 8192),
+            "temperature": payload.get('temperature', 0),
+            "top_p": payload.get('top_p', 0.7),
+            "messages": payload['messages'],
+            "stream": payload.get('stream', False),
+            "system": payload.get('system', "")
+        }
+    }
+
+    github_username, zed_user_id = random.choice(github_username_zed_userid_list)
+    jwt = create_jwt(github_username, zed_user_id)
+
+    headers = {
+        'Host': 'llm.zed.dev',
+        'accept': '*/*',
+        'content-type': 'application/json',
+        'authorization': f'Bearer {jwt}',
+        'user-agent': 'Zed/0.149.3 (macos; aarch64)'
+    }
+
+    # Get proxy from environment variable
+    proxy = os.environ.get('HTTP_PROXY', None)
+    proxies = {'http': proxy, 'https': proxy} if proxy else None
+
+    def generate():
+        with requests.post(url, headers=headers, json=llm_payload, stream=True, proxies=proxies) as response:
+            for chunk in response.iter_content(chunk_size=1024):
+                if chunk:
+                    # Parse the chunk and format it as per OpenAI's streaming format
+                    try:
+                        data = json.loads(chunk.decode('utf-8'))
+                        content = data.get('completion', '')
+                        yield f"data: {json.dumps({'choices': [{'delta': {'content': content}}]})}\n\n"
+                    except json.JSONDecodeError:
+                        continue
+            yield "data: [DONE]\n\n"
+
+    if payload.get('stream', False):
+        return Response(generate(), content_type='text/event-stream')
+    else:
+        with requests.post(url, headers=headers, json=llm_payload, proxies=proxies) as response:
+            data = response.json()
+            return jsonify({
+                "id": "chatcmpl-" + os.urandom(12).hex(),
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": model,
+                "choices": [{
+                    "index": 0,
+                    "message": {
+                        "role": "assistant",
+                        "content": data.get('completion', '')
+                    },
+                    "finish_reason": "stop"
+                }],
+                "usage": {
+                    "prompt_tokens": -1,  # We don't have this information
+                    "completion_tokens": -1,  # We don't have this information
+                    "total_tokens": -1  # We don't have this information
+                }
+            })
+
+if __name__ == '__main__':
+    app.run(debug=True)