Add ChatGPT-4 Model

Signed-off-by: Sebastiano Mariani <[email protected]>

Author: Phat3

pyproject.toml

@@ -47,3 +47,6 @@ ignored-modules = ""
 disable = """
 W1514,F0010,useless-super-delegation,E1103,W0108,W0404,R0904,R0922,W0105,
 W0142,C0301,C0321,C0322,C0324,R,W0232,E1001,W0212,W0703,C,I0011,I0012,I0013,E0012,W0511"""
+
+[project.scripts]
+llm-repl = "llm_repl.__main__:main"

src/llm_repl/__main__.py

@@ -11,7 +11,8 @@
     }
 }
 
-if __name__ == "__main__":
+
+def main():
     parser = argparse.ArgumentParser(description="LLM REPL")
     parser.add_argument(
         "--model",

src/llm_repl/llms/__init__.py

@@ -6,10 +6,13 @@
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.schema import AgentAction, AgentFinish, LLMResult
 
-from llm_repl.repl import LLMRepl
-
 
 class BaseLLM(ABC):
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Return the name of the LLM."""
+
     @property
     @abstractmethod
     def is_in_streaming_mode(self):
@@ -25,6 +28,9 @@ def process(self, msg) -> str:
         """Process the user message and return the response."""
 
 
+from llm_repl.repl import LLMRepl
+
+
 class StreamingCallbackHandler(BaseCallbackHandler):
     """Callback handler for streaming. Only works with LLMs that support streaming."""
 
@@ -80,11 +86,11 @@ def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
         """Run on agent end."""
 
 
-MODELS: Dict[str, Type[BaseLLM]] = {}
-
+from llm_repl.llms.chatgpt import ChatGPT
+from llm_repl.llms.chatgpt4 import ChatGPT4
 
-def register_model(model: Type[BaseLLM]):
-    """
-    Register a model in the list of available models.
-    """
-    MODELS[model.__name__.lower()] = model  # type: ignore
+# TODO: Implement dynamic loading of models
+MODELS: Dict[str, Type[BaseLLM]] = {
+    "chatgpt": ChatGPT,
+    "chatgpt4": ChatGPT4,
+}

src/llm_repl/llms/chatgpt.py

@@ -18,7 +18,7 @@
 from rich.markdown import Markdown
 
 from llm_repl.repl import LLMRepl
-from llm_repl.llms import BaseLLM, StreamingCallbackHandler, register_model
+from llm_repl.llms import BaseLLM, StreamingCallbackHandler  # , register_model
 
 
 class ChatGPTStreamingCallbackHandler(StreamingCallbackHandler):
@@ -54,7 +54,7 @@ def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
 
 
 class ChatGPT(BaseLLM):
-    def __init__(self, api_key: str, repl: LLMRepl):
+    def __init__(self, api_key: str, repl: LLMRepl, model_name: str = "gpt-3.5-turbo"):
         self.api_key = api_key
         # TODO: Make options configurable
         self.streaming_mode = True
@@ -76,11 +76,15 @@ def __init__(self, api_key: str, repl: LLMRepl):
             streaming=self.streaming_mode,
             callback_manager=CallbackManager([ChatGPTStreamingCallbackHandler(repl)]),
             verbose=True,
-            # temperature=0,
+            model_name=model_name,
         )  # type: ignore
         memory = ConversationBufferMemory(return_messages=True)
         self.model = ConversationChain(memory=memory, prompt=prompt, llm=llm)
 
+    @property
+    def name(self) -> str:
+        return "ChatGPT"
+
     @property
     def is_in_streaming_mode(self) -> bool:
         return self.streaming_mode
@@ -101,6 +105,3 @@ def load(cls, repl: LLMRepl) -> Optional[BaseLLM]:
     def process(self, msg: str) -> str:
         resp = self.model.predict(input=msg)
         return resp.strip()
-
-
-register_model(ChatGPT)

src/llm_repl/llms/chatgpt4.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import os
+
+from typing import Optional
+
+from llm_repl.repl import LLMRepl
+from llm_repl.llms import BaseLLM
+from llm_repl.llms.chatgpt import ChatGPT
+
+
+class ChatGPT4(ChatGPT):
+    @property
+    def name(self) -> str:
+        return "ChatGPT-4"
+
+    @classmethod
+    def load(cls, repl: LLMRepl) -> Optional[BaseLLM]:
+        api_key = os.getenv("OPENAI_API_KEY")
+        if api_key is None:
+            repl.print_error_msg(
+                "OpenAI API key not found. Please set the OPENAI_API_KEY environment variable."
+            )
+            return None
+
+        # TODO: Add autocomplete in repl
+        model = cls(api_key, repl, model_name="gpt-4")
+        return model

src/llm_repl/repl.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional, Type
+from typing import Any, Callable
 
 from prompt_toolkit import PromptSession
 from prompt_toolkit.completion import WordCompleter
@@ -7,6 +7,7 @@
 from rich.console import Console
 from rich.markdown import Markdown
 
+# from llm_repl.llms import BaseLLM
 
 LLM_CMD_HANDLERS: dict[str, Callable] = {}
 
@@ -34,7 +35,7 @@ def __init__(self, config: dict[str, Any]):
         self.server_color = config["style"]["server"]["color"]
         self.error_color = "bold red"
         self.misc_color = "gray"
-        self.model: Optional[BaseLLM] = None
+        self.model = None  # Optional[BaseLLM] = None
 
     def handle_enter(self, event):
         """
@@ -95,15 +96,15 @@ def print_error_msg(self, msg: str):
         """
         self._print_msg(self.ERROR_MSG_TITLE, msg, self.error_color)
 
-    def print_misc_msg(self, msg: str):
+    def print_misc_msg(self, msg: str, justify: str = "left"):
         """
         Print the miscellaneous message in the console.
 
         :param str msg: The message to be printed.
         """
-        self._print_msg("", msg, self.misc_color, justify="center")
+        self._print_msg("", msg, self.misc_color, justify=justify)
 
-    def run(self, model: Type[BaseLLM]):  # type: ignore
+    def run(self, model):
         """
         Starts the REPL.
 
@@ -115,11 +116,12 @@ def run(self, model: Type[BaseLLM]):  # type: ignore
         :param BaseLLM model: The LLM model to use.
         """
 
-        self.model = model.load(self)  # type: ignore
+        self.model = model.load(self)
         if self.model is None:
             return
 
-        self.print_misc_msg(self.INTRO_BANNER)
+        self.print_misc_msg(self.INTRO_BANNER, justify="center")
+        self.print_misc_msg(f"Loaded model: {self.model.name}", justify="center")
 
         while True:
             user_input = self.session.prompt("> ").rstrip()
@@ -140,6 +142,3 @@ def run(self, model: Type[BaseLLM]):  # type: ignore
             else:
                 self.console.print()
                 self.console.rule(style=self.server_color)
-
-
-from llm_repl.llms import BaseLLM