yangtb24 commited on
Commit
431b796
·
verified ·
1 Parent(s): c07387b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +252 -186
app.py CHANGED
@@ -546,9 +546,197 @@ def check_tokens():
546
 
547
  return jsonify(results)
548
 
549
- @app.route('/handsome/v1/chat/completions', methods=['POST'])
550
- def handsome_chat_completions():
551
- if not check_authorization(request):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
552
  return jsonify({"error": "Unauthorized"}), 401
553
 
554
  data = request.get_json()
@@ -656,7 +844,8 @@ def handsome_chat_completions():
656
  "index": 0,
657
  "message": {
658
  "role": "assistant",
659
- "content": image_url if image_url else "Failed to generate image", # Directly return the URL in content
 
660
  },
661
  "finish_reason": "length",
662
  }
@@ -667,7 +856,9 @@ def handsome_chat_completions():
667
  "total_tokens": 0
668
  }
669
  }
670
-
 
 
671
  except (KeyError, ValueError, IndexError) as e:
672
  logging.error(
673
  f"解析响应 JSON 失败: {e}, "
@@ -685,6 +876,7 @@ def handsome_chat_completions():
685
  "message": {
686
  "role": "assistant",
687
  "content": "Failed to process image data",
 
688
  },
689
  "finish_reason": "stop",
690
  }
@@ -702,10 +894,6 @@ def handsome_chat_completions():
702
  f"使用的模型: {model_name}"
703
  )
704
 
705
- with data_lock:
706
- request_timestamps.append(time.time())
707
- token_counts.append(0) # Image generation doesn't use tokens
708
-
709
  return jsonify(response_data)
710
  except requests.exceptions.RequestException as e:
711
  logging.error(f"请求转发异常: {e}")
@@ -718,191 +906,69 @@ def handsome_chat_completions():
718
  TEST_MODEL_ENDPOINT,
719
  headers=headers,
720
  json=data,
721
- stream=data.get("stream", False),
722
  timeout=60
723
  )
724
  if response.status_code == 429:
725
  return jsonify(response.json()), 429
726
-
727
- if data.get("stream", False):
728
- def generate():
729
- first_chunk_time = None
730
- full_response_content = ""
731
- for chunk in response.iter_content(chunk_size=1024):
732
- if chunk:
733
- if first_chunk_time is None:
734
- first_chunk_time = time.time()
735
- full_response_content += chunk.decode("utf-8")
736
- yield chunk
737
-
738
- end_time = time.time()
739
- first_token_time = (
740
- first_chunk_time - start_time
741
- if first_chunk_time else 0
742
- )
743
- total_time = end_time - start_time
744
-
745
- prompt_tokens = 0
746
- completion_tokens = 0
747
- response_content = ""
748
- for line in full_response_content.splitlines():
749
- if line.startswith("data:"):
750
- line = line[5:].strip()
751
- if line == "[DONE]":
752
- continue
753
- try:
754
- response_json = json.loads(line)
755
-
756
- if (
757
- "usage" in response_json and
758
- "completion_tokens" in response_json["usage"]
759
- ):
760
- completion_tokens = response_json[
761
- "usage"
762
- ]["completion_tokens"]
763
-
764
- if (
765
- "choices" in response_json and
766
- len(response_json["choices"]) > 0 and
767
- "delta" in response_json["choices"][0] and
768
- "content" in response_json[
769
- "choices"
770
- ][0]["delta"]
771
- ):
772
- response_content += response_json[
773
- "choices"
774
- ][0]["delta"]["content"]
775
-
776
- if (
777
- "usage" in response_json and
778
- "prompt_tokens" in response_json["usage"]
779
- ):
780
- prompt_tokens = response_json[
781
- "usage"
782
- ]["prompt_tokens"]
783
-
784
- except (
785
- KeyError,
786
- ValueError,
787
- IndexError
788
- ) as e:
789
- logging.error(
790
- f"解析流式响应单行 JSON 失败: {e}, "
791
- f"行内容: {line}"
792
  )
793
 
794
- user_content = ""
795
- messages = data.get("messages", [])
796
- for message in messages:
797
- if message["role"] == "user":
798
- if isinstance(message["content"], str):
799
- user_content += message["content"] + " "
800
- elif isinstance(message["content"], list):
801
- for item in message["content"]:
802
- if (
803
- isinstance(item, dict) and
804
- item.get("type") == "text"
805
- ):
806
- user_content += (
807
- item.get("text", "") +
808
- " "
809
- )
810
-
811
- user_content = user_content.strip()
812
-
813
- user_content_replaced = user_content.replace(
814
- '\n', '\\n'
815
- ).replace('\r', '\\n')
816
- response_content_replaced = response_content.replace(
817
- '\n', '\\n'
818
- ).replace('\r', '\\n')
819
-
820
- logging.info(
821
- f"使用的key: {api_key}, "
822
- f"提示token: {prompt_tokens}, "
823
- f"输出token: {completion_tokens}, "
824
- f"首字用时: {first_token_time:.4f}秒, "
825
- f"总共用时: {total_time:.4f}秒, "
826
- f"使用的模型: {model_name}, "
827
- f"用户的内容: {user_content_replaced}, "
828
- f"输出的内容: {response_content_replaced}"
829
- )
830
-
831
- with data_lock:
832
- request_timestamps.append(time.time())
833
- token_counts.append(prompt_tokens+completion_tokens)
834
 
835
- return Response(
836
- stream_with_context(generate()),
837
- content_type=response.headers['Content-Type']
838
- )
839
- else:
840
- response.raise_for_status()
841
- end_time = time.time()
842
- response_json = response.json()
843
- total_time = end_time - start_time
844
 
845
- try:
846
- prompt_tokens = response_json["usage"]["prompt_tokens"]
847
- completion_tokens = response_json[
848
- "usage"
849
- ]["completion_tokens"]
850
- response_content = response_json[
851
- "choices"
852
- ][0]["message"]["content"]
853
- except (KeyError, ValueError, IndexError) as e:
854
- logging.error(
855
- f"解析非流式响应 JSON 失败: {e}, "
856
- f"完整内容: {response_json}"
857
- )
858
- prompt_tokens = 0
859
- completion_tokens = 0
860
- response_content = ""
861
-
862
- user_content = ""
863
- messages = data.get("messages", [])
864
- for message in messages:
865
- if message["role"] == "user":
866
- if isinstance(message["content"], str):
867
- user_content += message["content"] + " "
868
- elif isinstance(message["content"], list):
869
- for item in message["content"]:
870
- if (
871
- isinstance(item, dict) and
872
- item.get("type") == "text"
873
- ):
874
- user_content += (
875
- item.get("text", "") +
876
- " "
877
- )
878
-
879
- user_content = user_content.strip()
880
-
881
- user_content_replaced = user_content.replace(
882
- '\n', '\\n'
883
- ).replace('\r', '\\n')
884
- response_content_replaced = response_content.replace(
885
- '\n', '\\n'
886
- ).replace('\r', '\\n')
887
-
888
- logging.info(
889
- f"使用的key: {api_key}, "
890
- f"提示token: {prompt_tokens}, "
891
- f"输出token: {completion_tokens}, "
892
- f"首字用时: 0, "
893
- f"总共用时: {total_time:.4f}秒, "
894
- f"使用的模型: {model_name}, "
895
- f"用户的内容: {user_content_replaced}, "
896
- f"输出的内容: {response_content_replaced}"
897
- )
898
- with data_lock:
899
- request_timestamps.append(time.time())
900
- if "prompt_tokens" in response_json["usage"] and "completion_tokens" in response_json["usage"]:
901
- token_counts.append(response_json["usage"]["prompt_tokens"] + response_json["usage"]["completion_tokens"])
902
- else:
903
- token_counts.append(0)
904
-
905
- return jsonify(response_json)
906
 
907
  except requests.exceptions.RequestException as e:
908
  logging.error(f"请求转发异常: {e}")
 
546
 
547
  return jsonify(results)
548
 
549
+ import os
550
+ import time
551
+ import logging
552
+ import requests
553
+ import json
554
+ import uuid
555
+ from flask import Flask, request, jsonify
556
+
557
+ API_ENDPOINT = "https://api.siliconflow.cn/v1/user/info"
558
+ TEST_MODEL_ENDPOINT = "https://api.siliconflow.cn/v1/chat/completions"
559
+ EMBEDDINGS_ENDPOINT = "https://api.siliconflow.cn/v1/embeddings"
560
+
561
+ text_models = []
562
+ free_text_models = []
563
+ embedding_models = []
564
+ free_embedding_models = []
565
+ image_models = []
566
+ free_image_models = []
567
+
568
+ invalid_keys_global = []
569
+ free_keys_global = []
570
+ unverified_keys_global = []
571
+ valid_keys_global = []
572
+ model_key_indices = {}
573
+
574
+ request_timestamps = []
575
+ token_counts = []
576
+ data_lock = None
577
+
578
+ FREE_MODEL_TEST_KEY = (
579
+ "sk-bmjbjzleaqfgtqfzmcnsbagxrlohriadnxqrzfocbizaxukw"
580
+ )
581
+
582
+ FREE_IMAGE_LIST = [
583
+ "stabilityai/stable-diffusion-3-5-large",
584
+ "black-forest-labs/FLUX.1-schnell",
585
+ "stabilityai/stable-diffusion-3-medium",
586
+ "stabilityai/stable-diffusion-xl-base-1.0",
587
+ "stabilityai/stable-diffusion-2-1"
588
+ ]
589
+
590
+ def get_credit_summary(api_key):
591
+ """
592
+ 使用 API 密钥获取额度信息。
593
+ """
594
+ headers = {
595
+ "Authorization": f"Bearer {api_key}",
596
+ "Content-Type": "application/json"
597
+ }
598
+ try:
599
+ response = requests.get(API_ENDPOINT, headers=headers)
600
+ response.raise_for_status()
601
+ data = response.json().get("data", {})
602
+ total_balance = data.get("totalBalance", 0)
603
+ return {"total_balance": float(total_balance)}
604
+ except requests.exceptions.RequestException as e:
605
+ logging.error(f"获取额度信息失败,API Key:{api_key},错误信息:{e}")
606
+ return None
607
+
608
+ def test_model_availability(api_key, model_name):
609
+ """
610
+ 测试指定的模型是否可用。
611
+ """
612
+ headers = {
613
+ "Authorization": f"Bearer {api_key}",
614
+ "Content-Type": "application/json"
615
+ }
616
+ try:
617
+ response = requests.post(
618
+ TEST_MODEL_ENDPOINT,
619
+ headers=headers,
620
+ json={
621
+ "model": model_name,
622
+ "messages": [{"role": "user", "content": "hi"}],
623
+ "max_tokens": 5,
624
+ "stream": False
625
+ },
626
+ timeout=5
627
+ )
628
+ if response.status_code == 429 or response.status_code == 200:
629
+ return True
630
+ else:
631
+ return False
632
+ except requests.exceptions.RequestException as e:
633
+ logging.error(
634
+ f"测试模型 {model_name} 可用性失败,"
635
+ f"API Key:{api_key},错误信息:{e}"
636
+ )
637
+ return False
638
+
639
+ def test_image_model_availability(api_key, model_name):
640
+ """
641
+ 测试指定的图像模型是否在 FREE_IMAGE_LIST 中。
642
+ 如果在列表中,返回 True,否则返回 False。
643
+ """
644
+ return model_name in FREE_IMAGE_LIST
645
+
646
+ def determine_request_type(model_name, model_list, free_model_list):
647
+ """
648
+ 根据用户请求的模型判断请求类型。
649
+ """
650
+ if model_name in free_model_list:
651
+ return "free"
652
+ elif model_name in model_list:
653
+ return "paid"
654
+ else:
655
+ return "unknown"
656
+
657
+ def select_key(request_type, model_name):
658
+ """
659
+ 根据请求类型和模型名称选择合适的 KEY,
660
+ 并实现轮询和重试机制。
661
+ """
662
+ if request_type == "free":
663
+ available_keys = (
664
+ free_keys_global +
665
+ unverified_keys_global +
666
+ valid_keys_global
667
+ )
668
+ elif request_type == "paid":
669
+ available_keys = unverified_keys_global + valid_keys_global
670
+ else:
671
+ available_keys = (
672
+ free_keys_global +
673
+ unverified_keys_global +
674
+ valid_keys_global
675
+ )
676
+
677
+ if not available_keys:
678
+ return None
679
+
680
+ current_index = model_key_indices.get(model_name, 0)
681
+
682
+ for _ in range(len(available_keys)):
683
+ key = available_keys[current_index % len(available_keys)]
684
+ current_index += 1
685
+
686
+ if key_is_valid(key, request_type):
687
+ model_key_indices[model_name] = current_index
688
+ return key
689
+ else:
690
+ logging.warning(
691
+ f"KEY {key} 无效或达到限制,尝试下一个 KEY"
692
+ )
693
+
694
+ model_key_indices[model_name] = 0
695
+ return None
696
+
697
+ def key_is_valid(key, request_type):
698
+ """
699
+ 检查 KEY 是否有效,
700
+ 根据不同的请求类型进行不同的检查。
701
+ """
702
+ if request_type == "invalid":
703
+ return False
704
+
705
+ credit_summary = get_credit_summary(key)
706
+ if credit_summary is None:
707
+ return False
708
+
709
+ total_balance = credit_summary.get("total_balance", 0)
710
+
711
+ if request_type == "free":
712
+ return True
713
+ elif request_type == "paid" or request_type == "unverified":
714
+ return total_balance > 0
715
+ else:
716
+ return False
717
+
718
+ def check_authorization(request, authorization_key):
719
+ """
720
+ 检查请求头中的 Authorization 字段
721
+ 是否匹配环境变量 AUTHORIZATION_KEY。
722
+ """
723
+ if not authorization_key:
724
+ logging.warning("环境变量 AUTHORIZATION_KEY 未设置,请设置后重试。")
725
+ return False
726
+
727
+ auth_header = request.headers.get('Authorization')
728
+ if not auth_header:
729
+ logging.warning("请求头中缺少 Authorization 字段。")
730
+ return False
731
+
732
+ if auth_header != f"Bearer {authorization_key}":
733
+ logging.warning(f"无效的 Authorization 密钥:{auth_header}")
734
+ return False
735
+
736
+ return True
737
+
738
+ def handsome_chat_completions(request, authorization_key):
739
+ if not check_authorization(request, authorization_key):
740
  return jsonify({"error": "Unauthorized"}), 401
741
 
742
  data = request.get_json()
 
844
  "index": 0,
845
  "message": {
846
  "role": "assistant",
847
+ "content": None,
848
+ "image_url": image_url if image_url else None, # Return URL in image_url
849
  },
850
  "finish_reason": "length",
851
  }
 
856
  "total_tokens": 0
857
  }
858
  }
859
+ if not image_url:
860
+ response_data["choices"][0]["message"]["content"] = "Failed to generate image"
861
+
862
  except (KeyError, ValueError, IndexError) as e:
863
  logging.error(
864
  f"解析响应 JSON 失败: {e}, "
 
876
  "message": {
877
  "role": "assistant",
878
  "content": "Failed to process image data",
879
+ "image_url": None,
880
  },
881
  "finish_reason": "stop",
882
  }
 
894
  f"使用的模型: {model_name}"
895
  )
896
 
 
 
 
 
897
  return jsonify(response_data)
898
  except requests.exceptions.RequestException as e:
899
  logging.error(f"请求转发异常: {e}")
 
906
  TEST_MODEL_ENDPOINT,
907
  headers=headers,
908
  json=data,
909
+ stream=False,
910
  timeout=60
911
  )
912
  if response.status_code == 429:
913
  return jsonify(response.json()), 429
914
+
915
+ response.raise_for_status()
916
+ end_time = time.time()
917
+ response_json = response.json()
918
+ total_time = end_time - start_time
919
+ try:
920
+ prompt_tokens = response_json["usage"]["prompt_tokens"]
921
+ completion_tokens = response_json[
922
+ "usage"
923
+ ]["completion_tokens"]
924
+ response_content = response_json[
925
+ "choices"
926
+ ][0]["message"]["content"]
927
+ except (KeyError, ValueError, IndexError) as e:
928
+ logging.error(
929
+ f"解析非流式响应 JSON 失败: {e}, "
930
+ f"完整内容: {response_json}"
931
+ )
932
+ prompt_tokens = 0
933
+ completion_tokens = 0
934
+ response_content = ""
935
+ user_content = ""
936
+ messages = data.get("messages", [])
937
+ for message in messages:
938
+ if message["role"] == "user":
939
+ if isinstance(message["content"], str):
940
+ user_content += message["content"] + " "
941
+ elif isinstance(message["content"], list):
942
+ for item in message["content"]:
943
+ if (
944
+ isinstance(item, dict) and
945
+ item.get("type") == "text"
946
+ ):
947
+ user_content += (
948
+ item.get("text", "") +
949
+ " "
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
950
  )
951
 
952
+ user_content = user_content.strip()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
953
 
954
+ user_content_replaced = user_content.replace(
955
+ '\n', '\\n'
956
+ ).replace('\r', '\\n')
957
+ response_content_replaced = response_content.replace(
958
+ '\n', '\\n'
959
+ ).replace('\r', '\\n')
 
 
 
960
 
961
+ logging.info(
962
+ f"使用的key: {api_key}, "
963
+ f"提示token: {prompt_tokens}, "
964
+ f"输出token: {completion_tokens}, "
965
+ f"首字用时: 0, "
966
+ f"总共用时: {total_time:.4f}秒, "
967
+ f"使用的模型: {model_name}, "
968
+ f"用户的内容: {user_content_replaced}, "
969
+ f"输出的内容: {response_content_replaced}"
970
+ )
971
+ return jsonify(response_json)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
972
 
973
  except requests.exceptions.RequestException as e:
974
  logging.error(f"请求转发异常: {e}")