iframe-proxy | Sunbelt Computer Software

77 lines (62 loc) · 2.16 KB

import ujson as json
import request
from usr.logging import getLogger
from usr.logging import getLogger
logger = getLogger(__name__)
class ARKConfig(object):
    CHAT_COMPLETIONS_POST_URL = 'https://ark.cn-beijing.volces.com/api/v3/chat/completions'
    MODEL_ID = 'xx'
    API_KEY = 'xx'
class ChatCompletionsError(Exception):
class ChatCompletions(object):
    def __init__(self, question):
        if not (isinstance(question, str) and question):
            raise ChatCompletionsError("question must be str type and not blank.")
        self.question = question
        self.resp = None
    def __enter__(self):
        self.__post()
        return self
    def __exit__(self, *args, **kwargs):
        if self.resp:
            self.resp.close()
    def __post(self):
        resp = request.post(
            ARKConfig.CHAT_COMPLETIONS_POST_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": "Bearer {}".format(ARKConfig.API_KEY)
            json={
                "model": ARKConfig.MODEL_ID,
                "messages": [
                    {"role": "user", "content": self.question},
                "stream": True
        if resp.status_code != 200:
            raise ChatCompletionsError("query_chat_completions exc: {}".format("".join((_ for _ in resp.text))))
        return resp
    @property
    def answer(self):
        self.resp = self.__post()
        raw = ""
        for temp in self.resp.text:
            raw += temp
            while True:
                data_index = raw.find("data: ")
                line_index = raw.find("\n\n")
                if data_index == -1 or line_index == -1:
                json_string = raw[data_index + 6 : line_index]
                if json_string == "[DONE]":
                data = json.loads(json_string)
                yield data["choices"][0]["delta"]["content"]
                raw = raw[line_index + 2:]
if __name__ == "__main__":
    with ChatCompletions("你好") as cc:
        for text in cc.answer:
            print(text, end="")

Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Sunbelt Computer Software

PL/B Language Development and Support

Uh oh!

FilesExpand file tree

ark_lib.py

Latest commit

History

ark_lib.py

File metadata and controls