Upload app.py

app.py

@@ -368,10 +368,10 @@ def handle_api_error(error, attempt):
                     logger.warning(f"用户输入因未知原因被阻止")
                     return 1, None
                 else:
-                    logger.warning(f"用户输入被阻止，原因未知: {block_reason_str}")
+                    logger.warning(f"用户输入被阻止，原因: {block_reason_str}")
                     return 1, None
             else:
-                logger.warning(f"用户输入被阻止，原因未知: {full_reason_str}")
+                logger.warning(f"用户输入被阻止，原因: {full_reason_str}")
                 return 1, None
 
         except (IndexError, AttributeError) as e:
@@ -585,6 +585,62 @@ def list_models():
     response = {"object": "list", "data": GEMINI_MODELS}
     return jsonify(response)
 
+@app.route('/hf/v1/embeddings', methods=['POST'])
+def embeddings():
+    data = request.get_json()
+    model_input = data.get("input", "text-embedding-004")
+    model = data.get("model")
+    if not model_input:
+        return jsonify({"error": "没有提供输入"}), 400
+
+    if isinstance(model_input, str):
+        model_input = [model_input]
+
+    gemini_request = {
+        "model": f"models/{model}",
+        "content": {
+            "parts": [{"text": text} for text in model_input]
+        }
+    }
+
+    gemini_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:embedContent?key={current_api_key}"
+    headers = {"Content-Type": "application/json"}
+    try:
+        gemini_response = requests.post(gemini_url, json=gemini_request, headers=headers)
+        gemini_response.raise_for_status()
+
+        response_json = gemini_response.json()
+        embeddings_data = []
+        if 'embedding' in response_json:
+          embeddings_data.append({
+              "object": "embedding",
+              "embedding": response_json['embedding']['values'],
+              "index": 0,
+          })
+        elif 'embeddings' in response_json:
+          for i, embedding in enumerate(response_json['embeddings']):
+              embeddings_data.append({
+                  "object": "embedding",
+                  "embedding": embedding['values'],
+                  "index": i,
+              })
+
+        client_response = {
+            "object": "list",
+            "data": embeddings_data,
+            "model": model,
+            "usage": {
+                "prompt_tokens": 0,
+                "total_tokens": 0,
+            },
+        }
+        switch_api_key()
+        return jsonify(client_response)
+
+    except requests.exceptions.RequestException as e:
+        print(f"请求Embeddings失败↙\: {e}")
+        return jsonify({"error": str(e)}), 500
+
 def keep_alive():
     try:
         response = requests.get("http://127.0.0.1:7860/", timeout=10)