知微
fa45d8aa5f
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
Privoxy对node122:18003返回500，直连正常
127 lines
3.4 KiB
Python
127 lines
3.4 KiB
Python
import json
|
||
|
||
import requests
|
||
try:
|
||
import sseclient
|
||
except:
|
||
raise ImportError("sseclient not found, please install it using 'pip install sseclient-py'.")
|
||
|
||
|
||
from tushare import get_token
|
||
|
||
BASE_URL = "http://api.waditu.com/dataapi"
|
||
# BASE_URL = "http://10.255.255.205:8083/dataapi"
|
||
API_KEY_PREFIX = "tsgpt-"
|
||
|
||
|
||
class GPTClient:
|
||
def __init__(self, token=None, timetout=120):
|
||
if not token:
|
||
token = get_token()
|
||
self.token = token
|
||
self.timeout = timetout
|
||
|
||
def _request(self, model, messages, temperature=None, max_tokens=None, stream=True, pretty=False) -> requests.Response:
|
||
"""
|
||
model string 模型名称， doubao-pro-128k
|
||
messages list 消息列表
|
||
[
|
||
{
|
||
"role": "user",
|
||
"content": "Hello World"
|
||
}
|
||
]
|
||
pretty bool 是否只返回回答内容文本
|
||
"""
|
||
resp = requests.post(
|
||
f'{BASE_URL}/llm/{model}',
|
||
json={"params": {
|
||
"stream": stream,
|
||
"messages": messages,
|
||
"temperature": temperature,
|
||
"max_tokens": max_tokens
|
||
}},
|
||
headers={"Authorization": f"tstoken-{self.token}"},
|
||
timeout=self.timeout, stream=stream
|
||
)
|
||
if resp.status_code != 200:
|
||
raise Exception(f"请求出现错误，{resp.content}")
|
||
return resp
|
||
|
||
def gpt_query(self, model, messages, temperature=None, max_tokens=None, pretty=False):
|
||
resp = self._request(model, messages, temperature, max_tokens, False, pretty)
|
||
resp_data = resp.json()
|
||
if resp_data.get('code') not in (0, None):
|
||
raise Exception(resp_data.get('msg') or resp_data)
|
||
if pretty:
|
||
return resp_data['choices'][0]["message"]["content"]
|
||
else:
|
||
return resp_data
|
||
|
||
def gpt_stream(self, model, messages, temperature=None, max_tokens=None, pretty=False):
|
||
resp = self._request(model, messages, temperature, max_tokens, True, pretty)
|
||
for e in sseclient.SSEClient(resp).events():
|
||
if '[DONE]' in e.data.upper():
|
||
break
|
||
e_data = json.loads(e.data)
|
||
if pretty:
|
||
yield e_data["choices"][0]["delta"]["content"]
|
||
else:
|
||
yield e_data
|
||
|
||
def gpt(self, model, query) -> str:
|
||
messages = [{
|
||
"role": "user",
|
||
"content": query
|
||
}]
|
||
return self.gpt_query(model, messages, pretty=True)
|
||
|
||
|
||
def test_gpt_query():
|
||
c = GPTClient()
|
||
dd = c.gpt_query("doubao-pro-128k", [{
|
||
"role": "user",
|
||
"content": "你好"
|
||
}])
|
||
print(dd)
|
||
dd = c.gpt_query("doubao-pro-128k", [{
|
||
"role": "user",
|
||
"content": "你好"
|
||
}], pretty=True)
|
||
print(dd)
|
||
|
||
|
||
def test_gpt_stream():
|
||
c = GPTClient()
|
||
# 流式接口， 返回每个数据对象
|
||
dd = c.gpt_stream("doubao-pro-128k", [
|
||
{
|
||
"role": "user",
|
||
"content": "你好"
|
||
}
|
||
])
|
||
for d in dd:
|
||
print(d)
|
||
|
||
# 流式接口， 返回每段文本，（pretty=True）
|
||
dd = c.gpt_stream("doubao-pro-128k", [
|
||
{
|
||
"role": "user",
|
||
"content": "你好"
|
||
}
|
||
], pretty=True)
|
||
for d in dd:
|
||
print(d)
|
||
|
||
|
||
def test_gpt():
|
||
# 非流式接口
|
||
c = GPTClient()
|
||
dd = c.gpt("doubao-pro-128k", "你好")
|
||
print(dd)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# test_gpt_stream()
|
||
test_gpt()
|