Files
MoFin/venv/lib/python3.12/site-packages/tushare/pro/llm.py
T
知微 fa45d8aa5f fix: 小果地址统一node122(兼容LAN+EasyTier)
- 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,直连正常
2026-06-30 02:56:35 +08:00

127 lines
3.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
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()