ipython-on-chatgpt.py

import os

import openai
import streamlit

openai.api_key = os.getenv("OPENAI_API_KEY")


@streamlit.cache_resource
def chat(messages) -> dict:
    result: dict = openai.ChatCompletion.create(  # type: ignore
        model="gpt-3.5-turbo",
        messages=messages,
    )
    return result


class ChatGPT:
    def __init__(self, system_prompt: str = ""):
        self.messages = []
        if system_prompt:
            self.messages = [
                {
                    "role": "system",
                    "content": system_prompt,
                }
            ]

    def add_sample(self, user_input: str, assistant_output: str):
        self.messages.append(
            {
                "role": "user",
                "content": user_input,
            }
        )
        self.messages.append(
            {
                "role": "assistant",
                "content": assistant_output,
            }
        )

    def chat(self, user_input: str) -> str:
        self.messages.append({"role": "user", "content": user_input})
        result: dict = chat(self.messages)
        content = result.get("choices", [])[0].get("message", {}).get("content")
        self.messages.append(
            {
                "role": "assistant",
                "content": content,
            }
        )
        return content


class Ipython(ChatGPT):
    def __init__(self, system_prompt: str = ""):
        self.time = 0
        super().__init__(system_prompt)

    def add_sample(self, user_input: str, assistant_output: str):
        user_content = f"In[{self.time}]: {user_input}"
        assistant_content = f"Out[{self.time}]: {assistant_output}"
        super().add_sample(user_content, assistant_content)
        self.time += 1

    def chat(self, user_input: str) -> str:
        user_content = f"In[{self.time}]: {user_input}"
        assistant_content = super().chat(user_content)
        self.time += 1
        return assistant_content


def run_ipython():
    client = Ipython()
    client.add_sample("'hello'", "'hello'")
    client.add_sample("1 + 2", "3")
    client.add_sample("2 * 3", "6")
    client.add_sample(
        """def f(x):
return 1 if x <= 1 else f(x-1) + f(x-2)

[f(x) for x in range(10)]
""",
        "[1, 1, 2, 3, 5, 8, 13, 21, 34, 55]",
    )
    client.add_sample(
        """client = ChatGPT(v=3.5)
client.ask("こんにちわ")
""",
        "'こんにちは！何かお話しましょうか？'",
    )
    client.add_sample(
        """@chatgpt.auto_implementation(prompt="アッカーマン関数を実装する")
def Ackermann(m: int, n: int) -> int:
  pass
print(Ackermann(3, 2))
""",
        "29",
    )

    for _ in range(2000):
        prompt = streamlit.text_area(label=f"In[{client.time}]")
        if prompt:
            content = client.chat(prompt)
            streamlit.code(content)
        else:
            break


run_ipython()