ernie_sdk.py

from dataclasses import asdict, dataclass
from typing import List

from pydantic import BaseModel, Field

from enum import Enum


class MessageRole(str, Enum):
    USER = "user"
    BOT = "assistant"


@dataclass
class Message:
    role: str
    content: str


@dataclass
class CompletionRequest:
    messages: List[Message]
    stream: bool = False
    user: str = ""


@dataclass
class Usage:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class CompletionResponse:
    id: str
    object: str
    created: int
    result: str
    need_clear_history: bool
    ban_round: int = 0
    sentence_id: int = 0
    is_end: bool = False
    usage: Usage = None
    is_safe: bool = False
    is_truncated: bool = False


class ErrorResponse(BaseModel):
    error_code: int = Field(...)
    error_msg: str = Field(...)
    id: str = Field(...)


class ErnieBot:
    url: str
    access_token: str
    request: CompletionRequest

    def __init__(self, url: str, access_token: str, request: CompletionRequest):
        self.url = url
        self.access_token = access_token
        self.request = request

    def get_response(self) -> CompletionResponse:

        import requests
        import json

        headers = {'Content-Type': 'application/json'}
        params = {'access_token': self.access_token}
        request_dict = asdict(self.request)
        response = requests.post(self.url, params=params, data=json.dumps(request_dict), headers=headers)
        # print(response.json())
        try:
            return CompletionResponse(**response.json())
        except Exception as e:
            print(e)
            raise Exception(response.json())


def user_message(prompt: str) -> Message:
    return Message(MessageRole.USER, prompt)


def bot_message(prompt: str) -> Message:
    return Message(MessageRole.BOT, prompt)