323 lines
12 KiB
Python
323 lines
12 KiB
Python
|
import contextlib
|
|||
|
|
import inspect
|
|||
|
|
import json
|
|||
|
|
import logging
|
|||
|
|
import os
|
|||
|
|
from typing import *
|
|||
|
|
|
|||
|
|
from open_chatcaht._constants import API_BASE_URI
|
|||
|
|
from open_chatcaht.utils import set_httpx_config, get_httpx_client, get_variable, get_function_default_params, \
|
|||
|
|
merge_dicts
|
|||
|
|
from functools import wraps
|
|||
|
|
from typing import Type, get_type_hints
|
|||
|
|
|
|||
|
|
import httpx
|
|||
|
|
import requests
|
|||
|
|
from pydantic import BaseModel
|
|||
|
|
|
|||
|
|
set_httpx_config()
|
|||
|
|
|
|||
|
|
CHATCHAT_API_BASE = get_variable(os.getenv('CHATCHAT_API_BASE'), 'http://127.0.0.1:8000')
|
|||
|
|
CHATCHAT_CLIENT_TIME_OUT = get_variable(os.getenv('CHATCHAT_CLIENT_TIME_OUT'), 60)
|
|||
|
|
CHATCHAT_CLIENT_DEFAULT_RETRY_COUNT = get_variable(os.getenv('CHATCHAT_CLIENT_DEFAULT_RETRY'), 3)
|
|||
|
|
CHATCHAT_CLIENT_DEFAULT_RETRY_INTERVAL = get_variable(os.getenv('CHATCHAT_CLIENT_DEFAULT_RETRY_INTERVAL'), 60)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ApiClient:
|
|||
|
|
"""
|
|||
|
|
api.py调用的封装（同步模式）,简化api调用方式
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(
|
|||
|
|
self,
|
|||
|
|
base_url: str = API_BASE_URI,
|
|||
|
|
timeout: float = 60,
|
|||
|
|
use_async: bool = False,
|
|||
|
|
use_proxy: bool = False,
|
|||
|
|
proxies=None,
|
|||
|
|
log_level: int = logging.INFO,
|
|||
|
|
retry: int = 3,
|
|||
|
|
retry_interval: int = 1,
|
|||
|
|
):
|
|||
|
|
if proxies is None:
|
|||
|
|
proxies = {}
|
|||
|
|
self.base_url = get_variable(base_url, CHATCHAT_API_BASE)
|
|||
|
|
self.timeout = get_variable(timeout, CHATCHAT_CLIENT_TIME_OUT)
|
|||
|
|
self._use_async = use_async
|
|||
|
|
self.use_proxy = use_proxy
|
|||
|
|
self.default_retry_count = get_variable(retry, CHATCHAT_CLIENT_DEFAULT_RETRY_COUNT)
|
|||
|
|
self.default_retry_interval = get_variable(retry_interval, CHATCHAT_CLIENT_DEFAULT_RETRY_INTERVAL)
|
|||
|
|
self.proxies = proxies
|
|||
|
|
self._client = None
|
|||
|
|
self.logger = logging.getLogger(__name__)
|
|||
|
|
self.logger.setLevel(log_level)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def client(self):
|
|||
|
|
if self._client is None or self._client.is_closed:
|
|||
|
|
self._client = get_httpx_client(
|
|||
|
|
base_url=self.base_url, use_async=self._use_async, timeout=self.timeout
|
|||
|
|
)
|
|||
|
|
return self._client
|
|||
|
|
|
|||
|
|
def _get(
|
|||
|
|
self,
|
|||
|
|
url: str,
|
|||
|
|
params: Union[Dict, List[Tuple], bytes] = None,
|
|||
|
|
retry: int = 3,
|
|||
|
|
stream: bool = False,
|
|||
|
|
**kwargs: Any,
|
|||
|
|
) -> Union[httpx.Response, Iterator[httpx.Response], None]:
|
|||
|
|
while retry > 0:
|
|||
|
|
try:
|
|||
|
|
if stream:
|
|||
|
|
return self.client.stream("GET", url, params=params, **kwargs)
|
|||
|
|
else:
|
|||
|
|
return self.client.get(url, params=params, **kwargs)
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"error when get {url}: {e}"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
retry -= 1
|
|||
|
|
|
|||
|
|
def _post(
|
|||
|
|
self,
|
|||
|
|
url: str,
|
|||
|
|
data: Dict = None,
|
|||
|
|
json: Dict = None,
|
|||
|
|
retry: int = 3,
|
|||
|
|
stream: bool = False,
|
|||
|
|
**kwargs: Any,
|
|||
|
|
) -> Union[httpx.Response, Iterator[httpx.Response], None]:
|
|||
|
|
while retry > 0:
|
|||
|
|
try:
|
|||
|
|
# print(kwargs)
|
|||
|
|
if stream:
|
|||
|
|
|
|||
|
|
return self.client.stream(
|
|||
|
|
"POST", url, data=data, json=json, **kwargs
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
self.logger.debug(f"post {url} with data: {data}")
|
|||
|
|
return self.client.post(url, data=data, json=json, **kwargs)
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"error when post {url}: {e}"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
retry -= 1
|
|||
|
|
|
|||
|
|
def _delete(
|
|||
|
|
self,
|
|||
|
|
url: str,
|
|||
|
|
data: Dict = None,
|
|||
|
|
json: Dict = None,
|
|||
|
|
retry: int = 3,
|
|||
|
|
stream: bool = False,
|
|||
|
|
**kwargs: Any,
|
|||
|
|
) -> Union[httpx.Response, Iterator[httpx.Response], None]:
|
|||
|
|
while retry > 0:
|
|||
|
|
try:
|
|||
|
|
if stream:
|
|||
|
|
return self.client.stream(
|
|||
|
|
"DELETE", url, data=data, json=json, **kwargs
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
return self.client.delete(url, data=data, json=json, **kwargs)
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"error when delete {url}: {e}"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
retry -= 1
|
|||
|
|
|
|||
|
|
def _httpx_stream2generator(
|
|||
|
|
self,
|
|||
|
|
response: contextlib._GeneratorContextManager,
|
|||
|
|
as_json: bool = False,
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
将httpx.stream返回的GeneratorContextManager转化为普通生成器
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
async def ret_async(response, as_json):
|
|||
|
|
try:
|
|||
|
|
async with response as r:
|
|||
|
|
chunk_cache = ""
|
|||
|
|
async for chunk in r.aiter_text(None):
|
|||
|
|
if not chunk: # fastchat api yield empty bytes on start and end
|
|||
|
|
continue
|
|||
|
|
if as_json:
|
|||
|
|
try:
|
|||
|
|
if chunk.startswith("data: "):
|
|||
|
|
data = json.loads(chunk_cache + chunk[6:-2])
|
|||
|
|
elif chunk.startswith(":"): # skip sse comment line
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
data = json.loads(chunk_cache + chunk)
|
|||
|
|
|
|||
|
|
chunk_cache = ""
|
|||
|
|
yield data
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"接口返回json错误： ‘{chunk}’。错误信息是：{e}。"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
|
|||
|
|
if chunk.startswith("data: "):
|
|||
|
|
chunk_cache += chunk[6:-2]
|
|||
|
|
elif chunk.startswith(":"): # skip sse comment line
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
chunk_cache += chunk
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
# print(chunk, end="", flush=True)
|
|||
|
|
yield chunk
|
|||
|
|
except httpx.ConnectError as e:
|
|||
|
|
msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。({e})"
|
|||
|
|
self.logger.error(msg)
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
except httpx.ReadTimeout as e:
|
|||
|
|
msg = f"API通信超时，请确认已启动FastChat与API服务（详见Wiki '5. 启动 API 服务或 Web UI'）。（{e}）"
|
|||
|
|
self.logger.error(msg)
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"API通信遇到错误：{e}"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
|
|||
|
|
def ret_sync(response, as_json):
|
|||
|
|
try:
|
|||
|
|
with response as r:
|
|||
|
|
chunk_cache = ""
|
|||
|
|
for chunk in r.iter_text(None):
|
|||
|
|
if not chunk: # fastchat api yield empty bytes on start and end
|
|||
|
|
continue
|
|||
|
|
if as_json:
|
|||
|
|
try:
|
|||
|
|
if chunk.startswith("data: "):
|
|||
|
|
data = json.loads(chunk_cache + chunk[6:-2])
|
|||
|
|
elif chunk.startswith(":"): # skip sse comment line
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
data = json.loads(chunk_cache + chunk)
|
|||
|
|
|
|||
|
|
chunk_cache = ""
|
|||
|
|
yield data
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"接口返回json错误： ‘{chunk}’。错误信息是：{e}。"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
|
|||
|
|
if chunk.startswith("data: "):
|
|||
|
|
chunk_cache += chunk[6:-2]
|
|||
|
|
elif chunk.startswith(":"): # skip sse comment line
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
chunk_cache += chunk
|
|||
|
|
continue
|
|||
|
|
else:
|
|||
|
|
# print(chunk, end="", flush=True)
|
|||
|
|
yield chunk
|
|||
|
|
except httpx.ConnectError as e:
|
|||
|
|
msg = f"无法连接API服务器，请确认 ‘api.py’ 已正常启动。({e})"
|
|||
|
|
self.logger.error(msg)
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
except httpx.ReadTimeout as e:
|
|||
|
|
msg = f"API通信超时，请确认已启动FastChat与API服务（详见Wiki '5. 启动 API 服务或 Web UI'）。（{e}）"
|
|||
|
|
self.logger.error(msg)
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = f"API通信遇到错误：{e}"
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
yield {"code": 500, "msg": msg}
|
|||
|
|
|
|||
|
|
if self._use_async:
|
|||
|
|
return ret_async(response, as_json)
|
|||
|
|
else:
|
|||
|
|
return ret_sync(response, as_json)
|
|||
|
|
|
|||
|
|
def _get_response_value(
|
|||
|
|
self,
|
|||
|
|
response: httpx.Response,
|
|||
|
|
as_json: bool = False,
|
|||
|
|
value_func: Callable = None,
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
转换同步或异步请求返回的响应
|
|||
|
|
`as_json`: 返回json
|
|||
|
|
`value_func`: 用户可以自定义返回值，该函数接受response或json
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def to_json(r):
|
|||
|
|
try:
|
|||
|
|
return r.json()
|
|||
|
|
except Exception as e:
|
|||
|
|
msg = "API未能返回正确的JSON。" + str(e)
|
|||
|
|
self.logger.error(f"{e.__class__.__name__}: {msg}")
|
|||
|
|
return {"code": 500, "msg": msg, "data": None}
|
|||
|
|
|
|||
|
|
if value_func is None:
|
|||
|
|
value_func = lambda r: r
|
|||
|
|
|
|||
|
|
async def ret_async(response):
|
|||
|
|
if as_json:
|
|||
|
|
return value_func(to_json(await response))
|
|||
|
|
else:
|
|||
|
|
return value_func(await response)
|
|||
|
|
|
|||
|
|
if self._use_async:
|
|||
|
|
return ret_async(response)
|
|||
|
|
else:
|
|||
|
|
if as_json:
|
|||
|
|
return value_func(to_json(response))
|
|||
|
|
else:
|
|||
|
|
return value_func(response)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_request_method(api_client_obj: ApiClient, method):
|
|||
|
|
if method is httpx.post:
|
|||
|
|
return getattr(api_client_obj, "_post")
|
|||
|
|
elif method is httpx.get:
|
|||
|
|
return getattr(api_client_obj, "_get")
|
|||
|
|
# elif method is httpx.put:
|
|||
|
|
# return api_client_obj.put
|
|||
|
|
elif method is httpx.delete:
|
|||
|
|
return getattr(api_client_obj, "_delete")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def http_request(method):
|
|||
|
|
def decorator(url, base_url='', headers=None, body_model: Type[BaseModel] = None, **options):
|
|||
|
|
headers = headers or {}
|
|||
|
|
|
|||
|
|
def wrapper(func):
|
|||
|
|
@wraps(func)
|
|||
|
|
def inner(*args, **kwargs):
|
|||
|
|
try:
|
|||
|
|
default_param: dict = get_function_default_params(func)
|
|||
|
|
|
|||
|
|
api_client_obj: ApiClient = args[0] if len(args) > 0 and isinstance(args[0], ApiClient) else None
|
|||
|
|
return_type = get_type_hints(func).get('return')
|
|||
|
|
full_url = base_url + url
|
|||
|
|
param = merge_dicts(kwargs, default_param)
|
|||
|
|
if body_model is not None:
|
|||
|
|
param = body_model(**kwargs).dict()
|
|||
|
|
# Send the HTTP request
|
|||
|
|
response = None
|
|||
|
|
if api_client_obj is not None:
|
|||
|
|
_method = get_request_method(api_client_obj, method)
|
|||
|
|
response = _method(full_url, headers=headers, json=param)
|
|||
|
|
else:
|
|||
|
|
response = method(full_url, headers=headers, json=param)
|
|||
|
|
response.raise_for_status()
|
|||
|
|
return response.json()
|
|||
|
|
except requests.exceptions.HTTPError as http_err:
|
|||
|
|
print(f"HTTP error occurred: {http_err}")
|
|||
|
|
except Exception as err:
|
|||
|
|
print(f"An error occurred: {err}")
|
|||
|
|
|
|||
|
|
return inner
|
|||
|
|
|
|||
|
|
return wrapper
|
|||
|
|
|
|||
|
|
return decorator
|
|||
|
|
|
|||
|
|
|
|||
|
|
post = http_request(httpx.post)
|
|||
|
|
get = http_request(httpx.get)
|
|||
|
|
delete = http_request(httpx.delete)
|
|||
|
|
put = http_request(httpx.put)
|