82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
import time
|
|
|
|
import requests
|
|
import json
|
|
from typing import List
|
|
|
|
from requests.adapters import HTTPAdapter
|
|
from urllib3 import Retry
|
|
from ..config.configs import config
|
|
import sseclient
|
|
|
|
|
|
class PromptItem:
|
|
def __init__(self, role: str, content: str):
|
|
self.role = role
|
|
self.content = content
|
|
|
|
def __str__(self):
|
|
return f"Promt[{self.role}:{self.content}]"
|
|
|
|
def bc_docker_ask_stream(
|
|
messages: List[PromptItem],
|
|
temperature=0.01,
|
|
repetition_penalty=1.05,
|
|
max_new_tokens=2048,
|
|
top_p=0.85,
|
|
top_k=5,
|
|
):
|
|
inputs = ""
|
|
for _ in messages:
|
|
inputs += "<C_Q>" if _.role == "user" else "<C_A>"
|
|
inputs += _.content
|
|
if len(messages) > 0 and messages[-1].role == "user":
|
|
inputs += "<C_A>"
|
|
params_bc_tgi = {
|
|
"inputs": inputs,
|
|
"parameters": {
|
|
"max_new_tokens": max_new_tokens,
|
|
"repetition_penalty": repetition_penalty,
|
|
"temperature": temperature,
|
|
"top_k": top_k,
|
|
"top_p": top_p,
|
|
"do_sample": True,
|
|
"details": False,
|
|
"stream": True,
|
|
}
|
|
}
|
|
result = None
|
|
retry_strategy = Retry(
|
|
total=1, # 最大重试次数（包括首次请求）
|
|
backoff_factor=1, # 重试之间的等待时间因子
|
|
status_forcelist=[429, 500, 502, 503, 504, 404], # 需要重试的状态码列表
|
|
allowed_methods=["POST"] # 只对POST请求进行重试
|
|
)
|
|
adapter = HTTPAdapter(max_retries=retry_strategy)
|
|
# 创建会话并添加重试逻辑
|
|
session = requests.Session()
|
|
session.mount("https://", adapter)
|
|
session.mount("http://", adapter)
|
|
for i in range(10):
|
|
try:
|
|
response = session.post(
|
|
config.get("bc_docker_api"), headers={"User-Agent": "bc-ccui-agents/0.0.0"},
|
|
json=params_bc_tgi, stream=True, timeout=30
|
|
)
|
|
if response.status_code != 200:
|
|
print(f"bc_docker_ask_stream: {response.status_code} {response.text}")
|
|
continue
|
|
sse = sseclient.SSEClient(response)
|
|
for msg in sse.events():
|
|
if msg.data != '[DONE]':
|
|
dd = json.loads(msg.data)['token'].get('text', '')
|
|
if dd.endswith("</s>"):
|
|
break
|
|
yield dd
|
|
break
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
time.sleep(2)
|
|
continue
|