| 37 | |
| 38 | |
| 39 | class ChatGPT: |
| 40 | |
| 41 | def __init__(self, messages: List[Dict] = None, model: str = "gpt-3.5-turbo", client: OpenAI = None, **model_kwargs): |
| 42 | """ |
| 43 | Create a chatgpt client |
| 44 | |
| 45 | :param messages: A list of messages comprising the conversation so far. |
| 46 | Each message is a dict with keys "role" and "content". |
| 47 | See: https://platform.openai.com/docs/api-reference/chat/create#chat/create-messages |
| 48 | :param model: The model to use. |
| 49 | :param OpenAI client: The openai client to use. If not provided, a new client will be created. |
| 50 | :param model_kwargs: Other parameters to pass to model, |
| 51 | See https://platform.openai.com/docs/api-reference/chat |
| 52 | """ |
| 53 | self._client = client or OpenAI() |
| 54 | self._messages = list(messages or []) |
| 55 | self.model_kwargs = dict(model=model, **model_kwargs) |
| 56 | |
| 57 | self.pending_stream_reply: ChatGPTStreamResponse = None |
| 58 | self.latest_nonstream_finish_reason = None |
| 59 | |
| 60 | def set_model(self, model: str): |
| 61 | """Set the model to use""" |
| 62 | self.model_kwargs['model'] = model |
| 63 | |
| 64 | def _ask(self, message: str, stream=True, **model_kwargs): |
| 65 | if self.pending_stream_reply: |
| 66 | self._messages.append({"role": "assistant", "content": self.pending_stream_reply.result()}) |
| 67 | self.pending_stream_reply = None |
| 68 | |
| 69 | self._messages.append({"role": "user", "content": message}) |
| 70 | resp = self._client.chat.completions.create( |
| 71 | **self.model_kwargs, |
| 72 | **model_kwargs, |
| 73 | messages=self._messages, |
| 74 | stream=stream, |
| 75 | ) |
| 76 | return resp |
| 77 | |
| 78 | def ask(self, message: str, **model_kwargs) -> str: |
| 79 | """ |
| 80 | Send a message to chatgpt and get the reply in string |
| 81 | |
| 82 | :param message: The message to send |
| 83 | :param model_kwargs: Other parameters to pass to openai.ChatCompletion.create() |
| 84 | :return: The reply from chatgpt |
| 85 | """ |
| 86 | resp = self._ask(message, stream=False, **model_kwargs) |
| 87 | reply = resp['choices'][0] |
| 88 | reply_content = reply['message']['content'] |
| 89 | self._messages.append({"role": "assistant", "content": reply_content}) |
| 90 | self.latest_nonstream_finish_reason = reply['finish_reason'] |
| 91 | |
| 92 | return reply_content |
| 93 | |
| 94 | def ask_stream(self, message: str, **model_kwargs) -> ChatGPTStreamResponse: |
| 95 | """ |
| 96 | Send a message to chatgpt and get the reply in stream |
no outgoing calls
no test coverage detected
searching dependent graphs…