95 lines
3.5 KiB
Python
95 lines
3.5 KiB
Python
from __future__ import annotations
|
||
|
||
from typing import Union, Mapping
|
||
|
||
from typing_extensions import override
|
||
|
||
from .core import _jwt_token
|
||
from .core import ZhipuAIError
|
||
from .core import HttpClient, ZHIPUAI_DEFAULT_MAX_RETRIES
|
||
from .core import NotGiven, NOT_GIVEN
|
||
from . import api_resource
|
||
import os
|
||
import httpx
|
||
from httpx import Timeout
|
||
|
||
|
||
class ZhipuAI(HttpClient):
|
||
chat: api_resource.chat.Chat
|
||
api_key: str
|
||
_disable_token_cache: bool = True
|
||
source_channel: str
|
||
|
||
def __init__(
|
||
self,
|
||
*,
|
||
api_key: str | None = None,
|
||
base_url: str | httpx.URL | None = None,
|
||
timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
|
||
max_retries: int = ZHIPUAI_DEFAULT_MAX_RETRIES,
|
||
http_client: httpx.Client | None = None,
|
||
custom_headers: Mapping[str, str] | None = None,
|
||
disable_token_cache: bool = True,
|
||
_strict_response_validation: bool = False,
|
||
source_channel: str | None = None
|
||
) -> None:
|
||
if api_key is None:
|
||
api_key = os.environ.get("ZHIPUAI_API_KEY")
|
||
if api_key is None:
|
||
raise ZhipuAIError("未提供api_key，请通过参数或环境变量提供")
|
||
self.api_key = api_key
|
||
self.source_channel = source_channel
|
||
self._disable_token_cache = disable_token_cache
|
||
|
||
if base_url is None:
|
||
base_url = os.environ.get("ZHIPUAI_BASE_URL")
|
||
if base_url is None:
|
||
base_url = f"https://open.bigmodel.cn/api/paas/v4"
|
||
from .__version__ import __version__
|
||
super().__init__(
|
||
version=__version__,
|
||
base_url=base_url,
|
||
max_retries=max_retries,
|
||
timeout=timeout,
|
||
custom_httpx_client=http_client,
|
||
custom_headers=custom_headers,
|
||
_strict_response_validation=_strict_response_validation,
|
||
)
|
||
self.chat = api_resource.chat.Chat(self)
|
||
self.images = api_resource.images.Images(self)
|
||
self.embeddings = api_resource.embeddings.Embeddings(self)
|
||
self.files = api_resource.files.Files(self)
|
||
self.fine_tuning = api_resource.fine_tuning.FineTuning(self)
|
||
self.batches = api_resource.Batches(self)
|
||
self.knowledge = api_resource.Knowledge(self)
|
||
self.tools = api_resource.Tools(self)
|
||
self.videos = api_resource.Videos(self)
|
||
self.assistant = api_resource.Assistant(self)
|
||
self.web_search = api_resource.WebSearchApi(self)
|
||
self.audio = api_resource.audio.Audio(self)
|
||
self.moderations = api_resource.moderation.Moderations(self)
|
||
self.agents = api_resource.agents.Agents(self)
|
||
self.file_parser = api_resource.file_parser.FileParser(self)
|
||
|
||
@property
|
||
@override
|
||
def auth_headers(self) -> dict[str, str]:
|
||
api_key = self.api_key
|
||
source_channel = self.source_channel or "python-sdk"
|
||
if self._disable_token_cache:
|
||
return {"Authorization": f"Bearer {api_key}","x-source-channel": source_channel}
|
||
else:
|
||
return {"Authorization": f"Bearer {_jwt_token.generate_token(api_key)}","x-source-channel": source_channel}
|
||
|
||
def __del__(self) -> None:
|
||
if (not hasattr(self, "_has_custom_http_client")
|
||
or not hasattr(self, "close")
|
||
or not hasattr(self, "_client")):
|
||
# if the '__init__' method raised an error, self would not have client attr
|
||
return
|
||
|
||
if self._has_custom_http_client:
|
||
return
|
||
|
||
self.close()
|