PyWebIO/demos/chatgpt.py

import json
import time
from typing import Dict, List

from openai import OpenAI, Stream
from openai.types.chat import ChatCompletionChunk

import pywebio_battery
from pywebio.input import *
from pywebio.output import *
from pywebio.pin import *
from pywebio.session import set_env, download


class ChatGPTStreamResponse:
    """
    A wrapper to Stream[ChatCompletionChunk], add a `result()` method to get the final result.
    """
    def __init__(self, response: Stream[ChatCompletionChunk]):
        self.response = response
        self.yielded = []
        self.finish_reason = None

    def __next__(self):
        chunk = next(self.response)
        self.finish_reason = chunk.choices[0].finish_reason
        delta = chunk.choices[0].delta
        if delta.content:
            self.yielded.append(delta.content)
        return delta.content

    def __iter__(self):
        return self

    def result(self):
        return ''.join(self.yielded)


class ChatGPT:

    def __init__(self, messages: List[Dict] = None, model: str = "gpt-3.5-turbo", client: OpenAI = None, **model_kwargs):
        """
        Create a chatgpt client

        :param messages: A list of messages comprising the conversation so far.
          Each message is a dict with keys "role" and "content".
          See: https://platform.openai.com/docs/api-reference/chat/create#chat/create-messages
        :param model: The model to use.
        :param OpenAI client: The openai client to use. If not provided, a new client will be created.
        :param model_kwargs: Other parameters to pass to model,
          See https://platform.openai.com/docs/api-reference/chat
        """
        self._client = client or OpenAI()
        self._messages = list(messages or [])
        self.model_kwargs = dict(model=model, **model_kwargs)

        self.pending_stream_reply: ChatGPTStreamResponse = None
        self.latest_nonstream_finish_reason = None

    def set_model(self, model: str):
        """Set the model to use"""
        self.model_kwargs['model'] = model

    def _ask(self, message: str, stream=True, **model_kwargs):
        if self.pending_stream_reply:
            self._messages.append({"role": "assistant", "content": self.pending_stream_reply.result()})
            self.pending_stream_reply = None

        self._messages.append({"role": "user", "content": message})
        resp = self._client.chat.completions.create(
            **self.model_kwargs,
            **model_kwargs,
            messages=self._messages,
            stream=stream,
        )
        return resp

    def ask(self, message: str, **model_kwargs) -> str:
        """
        Send a message to chatgpt and get the reply in string

        :param message: The message to send
        :param model_kwargs: Other parameters to pass to openai.ChatCompletion.create()
        :return: The reply from chatgpt
        """
        resp = self._ask(message, stream=False, **model_kwargs)
        reply = resp['choices'][0]
        reply_content = reply['message']['content']
        self._messages.append({"role": "assistant", "content": reply_content})
        self.latest_nonstream_finish_reason = reply['finish_reason']

        return reply_content

    def ask_stream(self, message: str, **model_kwargs) -> ChatGPTStreamResponse:
        """
        Send a message to chatgpt and get the reply in stream

        :param message: The message to send
        :param model_kwargs: Other parameters to pass to openai.ChatCompletion.create()
        :return: A iterator that yields the reply from chatgpt.
          The iterator will be exhausted when the reply is complete.
        """
        resp = self._ask(message, stream=True, **model_kwargs)
        self.pending_stream_reply = ChatGPTStreamResponse(resp)
        return self.pending_stream_reply

    def latest_finish_reason(self) -> str:
        """The finish reason for the latest reply of chatgpt.

        The possible values for finish_reason are:
          'stop': API returned complete model output
          'length': Incomplete model output due to max_tokens parameter or token limit
          'content_filter': Omitted content due to a flag from our content filters
          'null': API response still in progress or incomplete

        See: https://platform.openai.com/docs/guides/chat/response-format
        """
        if self.pending_stream_reply:
            return self.pending_stream_reply.finish_reason
        return self.latest_nonstream_finish_reason

    def messages(self) -> List[Dict]:
        """Get all messages of the conversation """
        if self.pending_stream_reply:
            self._messages.append({"role": "assistant", "content": self.pending_stream_reply.result()})
            self.pending_stream_reply = None

        return self._messages


def get_openai_config():
    openai_config = json.loads(pywebio_battery.get_localstorage('openai_config') or '{}')
    if not openai_config:
        openai_config = input_group('OpenAI API Config', [
            input('API Key', name='api_key', type=TEXT, required=True,
                  help_text='Get your API key from https://platform.openai.com/account/api-keys'),
            input('API Server', name='api_base', type=TEXT, value='https://api.openai.com', required=True),
        ])
        openai_config['api_base'] = openai_config['api_base'].removesuffix('/v1').strip('/') + '/v1'
        pywebio_battery.set_localstorage('openai_config', json.dumps(openai_config))

    put_button('Reset OpenAI API Key', reset_openai_config, link_style=True)
    return openai_config


def reset_openai_config():
    pywebio_battery.set_localstorage('openai_config', json.dumps(None))
    toast("Please refresh the page to take effect")


def main():
    """"""
    set_env(input_panel_fixed=False, output_animation=False)
    put_markdown("""
    # ChatGPT
    A ChatGPT client implemented with PyWebIO. [Source Code](https://github.com/pywebio/PyWebIO/blob/dev/demos/chatgpt.py)
    TIPS: refresh page to open a new chat.
    """)
    put_select('model', ['gpt-3.5-turbo', 'gpt-4'], label='Model')

    openai_config = get_openai_config()
    client = OpenAI(api_key=openai_config['api_key'], base_url=openai_config['api_base'])

    bot = ChatGPT(client=client, model=pin.model)
    pin_on_change('model', lambda v: bot.set_model(v))
    while True:
        form = input_group('', [
            input(name='msg', placeholder='Ask ChatGPT'),
            actions(name='cmd', buttons=['Send', 'Multi-line Input', 'Save Chat'])
        ])
        if form['cmd'] == 'Multi-line Input':
            form['msg'] = textarea(value=form['msg'])
        elif form['cmd'] == 'Save Chat':
            messages = [
                msg['content'] if msg['role'] == 'user' else f"> {msg['content']}"
                for msg in bot.messages()
            ]
            download(f"chatgpt_{time.strftime('%Y%m%d%H%M%S')}.md",
                     '\n\n'.join(messages).encode('utf8'))
            continue

        user_msg = form['msg']
        if not user_msg:
            continue

        put_info(put_text(user_msg, inline=True))

        with use_scope(f'reply-{int(time.time())}'):
            put_loading('grow', 'info')
            try:
                reply_chunks = bot.ask_stream(user_msg)
            except Exception as e:
                popup('ChatGPT Error', put_error(e))
                continue
            finally:
                clear()  # clear loading
            for chunk in reply_chunks:
                put_text(chunk, inline=True)
            clear()  # clear above text
            put_markdown(reply_chunks.result())

        if bot.latest_finish_reason() == 'length':
            put_error('Incomplete model output due to max_tokens parameter or token limit.')
        elif bot.latest_finish_reason() == 'content_filter':
            put_warning("Omitted content due to a flag from OpanAI's content filters.")


if __name__ == '__main__':
    from pywebio import start_server

    start_server(main, port=8080, debug=True, cdn=False)