2
2
import time
3
3
from typing import Dict , List
4
4
5
- import openai
5
+ from openai import OpenAI , Stream
6
+ from openai .types .chat import ChatCompletionChunk
6
7
7
8
import pywebio_battery
8
9
from pywebio .input import *
12
13
13
14
14
15
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 ]):
16
20
self .response = response
17
21
self .yielded = []
18
22
self .finish_reason = None
19
23
20
24
def __next__ (self ):
21
- # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_stream_completions.ipynb
22
25
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
31
31
32
32
def __iter__ (self ):
33
33
return self
@@ -38,35 +38,36 @@ def result(self):
38
38
39
39
class ChatGPT :
40
40
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 ):
42
42
"""
43
43
Create a chatgpt client
44
44
45
45
:param messages: A list of messages comprising the conversation so far.
46
46
Each message is a dict with keys "role" and "content".
47
47
See: https://platform.openai.com/docs/api-reference/chat/create#chat/create-messages
48
48
: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.
51
50
:param model_kwargs: Other parameters to pass to model,
52
51
See https://platform.openai.com/docs/api-reference/chat
53
52
"""
53
+ self ._client = client or OpenAI ()
54
54
self ._messages = list (messages or [])
55
55
self .model_kwargs = dict (model = model , ** model_kwargs )
56
- if api_key :
57
- self .model_kwargs ['api_key' ] = api_key
58
56
59
57
self .pending_stream_reply : ChatGPTStreamResponse = None
60
58
self .latest_nonstream_finish_reason = None
61
59
60
+ def set_model (self , model : str ):
61
+ """Set the model to use"""
62
+ self .model_kwargs ['model' ] = model
63
+
62
64
def _ask (self , message : str , stream = True , ** model_kwargs ):
63
65
if self .pending_stream_reply :
64
66
self ._messages .append ({"role" : "assistant" , "content" : self .pending_stream_reply .result ()})
65
67
self .pending_stream_reply = None
66
68
67
69
self ._messages .append ({"role" : "user" , "content" : message })
68
-
69
- resp = openai .ChatCompletion .create (
70
+ resp = self ._client .chat .completions .create (
70
71
** self .model_kwargs ,
71
72
** model_kwargs ,
72
73
messages = self ._messages ,
@@ -158,8 +159,10 @@ def main():
158
159
put_select ('model' , ['gpt-3.5-turbo' , 'gpt-4' ], label = 'Model' )
159
160
160
161
openai_config = get_openai_config ()
162
+ client = OpenAI (api_key = openai_config ['api_key' ], base_url = openai_config ['api_base' ])
161
163
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 ))
163
166
while True :
164
167
form = input_group ('' , [
165
168
input (name = 'msg' , placeholder = 'Ask ChatGPT' ),
0 commit comments