Skip to content

Commit 51ece4a

Browse files
committed
update chatgpt demo
1 parent bdd6c83 commit 51ece4a

File tree

1 file changed

+22
-19
lines changed

1 file changed

+22
-19
lines changed

demos/chatgpt.py

Lines changed: 22 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
import time
33
from typing import Dict, List
44

5-
import openai
5+
from openai import OpenAI, Stream
6+
from openai.types.chat import ChatCompletionChunk
67

78
import pywebio_battery
89
from pywebio.input import *
@@ -12,22 +13,21 @@
1213

1314

1415
class ChatGPTStreamResponse:
15-
def __init__(self, response):
16+
"""
17+
A wrapper to Stream[ChatCompletionChunk], add a `result()` method to get the final result.
18+
"""
19+
def __init__(self, response: Stream[ChatCompletionChunk]):
1620
self.response = response
1721
self.yielded = []
1822
self.finish_reason = None
1923

2024
def __next__(self):
21-
# https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
2225
chunk = next(self.response)
23-
self.finish_reason = chunk['choices'][0]['finish_reason']
24-
25-
# { "role": "assistant" } or { "content": "..."} or {}
26-
delta = chunk['choices'][0]['delta']
27-
content = delta.get('content', '')
28-
if content:
29-
self.yielded.append(content)
30-
return content
26+
self.finish_reason = chunk.choices[0].finish_reason
27+
delta = chunk.choices[0].delta
28+
if delta.content:
29+
self.yielded.append(delta.content)
30+
return delta.content
3131

3232
def __iter__(self):
3333
return self
@@ -38,35 +38,36 @@ def result(self):
3838

3939
class ChatGPT:
4040

41-
def __init__(self, messages: List[Dict] = None, model: str = "gpt-3.5-turbo", api_key=None, **model_kwargs):
41+
def __init__(self, messages: List[Dict] = None, model: str = "gpt-3.5-turbo", client: OpenAI = None, **model_kwargs):
4242
"""
4343
Create a chatgpt client
4444
4545
:param messages: A list of messages comprising the conversation so far.
4646
Each message is a dict with keys "role" and "content".
4747
See: https://platform.openai.com/docs/api-reference/chat/create#chat/create-messages
4848
:param model: The model to use.
49-
:param api_key: The openai api key.
50-
Get your API key from https://platform.openai.com/account/api-keys
49+
:param OpenAI client: The openai client to use. If not provided, a new client will be created.
5150
:param model_kwargs: Other parameters to pass to model,
5251
See https://platform.openai.com/docs/api-reference/chat
5352
"""
53+
self._client = client or OpenAI()
5454
self._messages = list(messages or [])
5555
self.model_kwargs = dict(model=model, **model_kwargs)
56-
if api_key:
57-
self.model_kwargs['api_key'] = api_key
5856

5957
self.pending_stream_reply: ChatGPTStreamResponse = None
6058
self.latest_nonstream_finish_reason = None
6159

60+
def set_model(self, model: str):
61+
"""Set the model to use"""
62+
self.model_kwargs['model'] = model
63+
6264
def _ask(self, message: str, stream=True, **model_kwargs):
6365
if self.pending_stream_reply:
6466
self._messages.append({"role": "assistant", "content": self.pending_stream_reply.result()})
6567
self.pending_stream_reply = None
6668

6769
self._messages.append({"role": "user", "content": message})
68-
69-
resp = openai.ChatCompletion.create(
70+
resp = self._client.chat.completions.create(
7071
**self.model_kwargs,
7172
**model_kwargs,
7273
messages=self._messages,
@@ -158,8 +159,10 @@ def main():
158159
put_select('model', ['gpt-3.5-turbo', 'gpt-4'], label='Model')
159160

160161
openai_config = get_openai_config()
162+
client = OpenAI(api_key=openai_config['api_key'], base_url=openai_config['api_base'])
161163

162-
bot = ChatGPT(api_key=openai_config['api_key'], api_base=openai_config['api_base'], model=pin.model)
164+
bot = ChatGPT(client=client, model=pin.model)
165+
pin_on_change('model', lambda v: bot.set_model(v))
163166
while True:
164167
form = input_group('', [
165168
input(name='msg', placeholder='Ask ChatGPT'),

0 commit comments

Comments
 (0)