2024-04-05 19:05:00 -04:00
|
|
|
from __future__ import annotations
|
|
|
|
|
|
|
|
import json
|
|
|
|
|
2024-04-10 02:14:50 -04:00
|
|
|
from ..helper import filter_none
|
2024-04-05 19:05:00 -04:00
|
|
|
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin, FinishReason
|
2024-05-05 17:38:31 -04:00
|
|
|
from ...typing import Union, Optional, AsyncResult, Messages, ImageType
|
2024-04-10 02:14:50 -04:00
|
|
|
from ...requests import StreamSession, raise_for_status
|
2024-04-07 04:36:13 -04:00
|
|
|
from ...errors import MissingAuthError, ResponseError
|
2024-05-05 17:38:31 -04:00
|
|
|
from ...image import to_data_uri
|
2024-04-05 19:05:00 -04:00
|
|
|
|
|
|
|
class Openai(AsyncGeneratorProvider, ProviderModelMixin):
|
2024-04-12 14:58:40 -04:00
|
|
|
label = "OpenAI API"
|
2024-04-05 19:05:00 -04:00
|
|
|
url = "https://openai.com"
|
|
|
|
working = True
|
|
|
|
needs_auth = True
|
|
|
|
supports_message_history = True
|
|
|
|
supports_system_message = True
|
|
|
|
|
|
|
|
@classmethod
|
|
|
|
async def create_async_generator(
|
|
|
|
cls,
|
|
|
|
model: str,
|
|
|
|
messages: Messages,
|
|
|
|
proxy: str = None,
|
|
|
|
timeout: int = 120,
|
2024-05-05 17:38:31 -04:00
|
|
|
image: ImageType = None,
|
2024-04-05 19:05:00 -04:00
|
|
|
api_key: str = None,
|
|
|
|
api_base: str = "https://api.openai.com/v1",
|
|
|
|
temperature: float = None,
|
|
|
|
max_tokens: int = None,
|
|
|
|
top_p: float = None,
|
2024-04-07 04:36:13 -04:00
|
|
|
stop: Union[str, list[str]] = None,
|
2024-04-05 19:05:00 -04:00
|
|
|
stream: bool = False,
|
2024-04-07 04:36:13 -04:00
|
|
|
headers: dict = None,
|
|
|
|
extra_data: dict = {},
|
2024-04-05 19:05:00 -04:00
|
|
|
**kwargs
|
|
|
|
) -> AsyncResult:
|
2024-04-07 04:36:13 -04:00
|
|
|
if cls.needs_auth and api_key is None:
|
2024-04-05 19:05:00 -04:00
|
|
|
raise MissingAuthError('Add a "api_key"')
|
2024-05-05 17:38:31 -04:00
|
|
|
if image is not None:
|
|
|
|
if not model and hasattr(cls, "default_vision_model"):
|
|
|
|
model = cls.default_vision_model
|
|
|
|
messages[-1]["content"] = [
|
|
|
|
{
|
|
|
|
"type": "image_url",
|
|
|
|
"image_url": {"url": to_data_uri(image)}
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"type": "text",
|
|
|
|
"text": messages[-1]["content"]
|
|
|
|
}
|
|
|
|
]
|
2024-04-05 19:05:00 -04:00
|
|
|
async with StreamSession(
|
|
|
|
proxies={"all": proxy},
|
2024-04-07 04:36:13 -04:00
|
|
|
headers=cls.get_headers(stream, api_key, headers),
|
2024-04-05 19:05:00 -04:00
|
|
|
timeout=timeout
|
|
|
|
) as session:
|
2024-04-07 04:36:13 -04:00
|
|
|
data = filter_none(
|
|
|
|
messages=messages,
|
|
|
|
model=cls.get_model(model),
|
|
|
|
temperature=temperature,
|
|
|
|
max_tokens=max_tokens,
|
|
|
|
top_p=top_p,
|
|
|
|
stop=stop,
|
|
|
|
stream=stream,
|
|
|
|
**extra_data
|
|
|
|
)
|
2024-04-05 19:05:00 -04:00
|
|
|
async with session.post(f"{api_base.rstrip('/')}/chat/completions", json=data) as response:
|
|
|
|
await raise_for_status(response)
|
2024-04-07 04:36:13 -04:00
|
|
|
if not stream:
|
|
|
|
data = await response.json()
|
2024-04-18 14:18:51 -04:00
|
|
|
cls.raise_error(data)
|
2024-04-07 04:36:13 -04:00
|
|
|
choice = data["choices"][0]
|
|
|
|
if "content" in choice["message"]:
|
|
|
|
yield choice["message"]["content"].strip()
|
|
|
|
finish = cls.read_finish_reason(choice)
|
|
|
|
if finish is not None:
|
|
|
|
yield finish
|
|
|
|
else:
|
|
|
|
first = True
|
|
|
|
async for line in response.iter_lines():
|
|
|
|
if line.startswith(b"data: "):
|
|
|
|
chunk = line[6:]
|
|
|
|
if chunk == b"[DONE]":
|
|
|
|
break
|
|
|
|
data = json.loads(chunk)
|
2024-04-18 14:18:51 -04:00
|
|
|
cls.raise_error(data)
|
2024-04-07 04:36:13 -04:00
|
|
|
choice = data["choices"][0]
|
|
|
|
if "content" in choice["delta"] and choice["delta"]["content"]:
|
|
|
|
delta = choice["delta"]["content"]
|
|
|
|
if first:
|
|
|
|
delta = delta.lstrip()
|
|
|
|
if delta:
|
|
|
|
first = False
|
|
|
|
yield delta
|
|
|
|
finish = cls.read_finish_reason(choice)
|
|
|
|
if finish is not None:
|
|
|
|
yield finish
|
2024-04-05 19:05:00 -04:00
|
|
|
|
|
|
|
@staticmethod
|
2024-04-07 04:36:13 -04:00
|
|
|
def read_finish_reason(choice: dict) -> Optional[FinishReason]:
|
2024-04-05 19:05:00 -04:00
|
|
|
if "finish_reason" in choice and choice["finish_reason"] is not None:
|
2024-04-07 04:36:13 -04:00
|
|
|
return FinishReason(choice["finish_reason"])
|
2024-04-05 19:05:00 -04:00
|
|
|
|
2024-04-18 14:18:51 -04:00
|
|
|
@staticmethod
|
|
|
|
def raise_error(data: dict):
|
|
|
|
if "error_message" in data:
|
|
|
|
raise ResponseError(data["error_message"])
|
|
|
|
elif "error" in data:
|
|
|
|
raise ResponseError(f'Error {data["error"]["code"]}: {data["error"]["message"]}')
|
|
|
|
|
2024-04-07 04:36:13 -04:00
|
|
|
@classmethod
|
|
|
|
def get_headers(cls, stream: bool, api_key: str = None, headers: dict = None) -> dict:
|
2024-04-05 19:05:00 -04:00
|
|
|
return {
|
2024-04-07 04:36:13 -04:00
|
|
|
"Accept": "text/event-stream" if stream else "application/json",
|
2024-04-05 19:05:00 -04:00
|
|
|
"Content-Type": "application/json",
|
2024-04-07 04:36:13 -04:00
|
|
|
**(
|
|
|
|
{"Authorization": f"Bearer {api_key}"}
|
2024-05-05 17:38:31 -04:00
|
|
|
if api_key is not None else {}
|
2024-04-07 04:36:13 -04:00
|
|
|
),
|
|
|
|
**({} if headers is None else headers)
|
2024-04-10 02:14:50 -04:00
|
|
|
}
|