Langchain integration (AntonOsika#512)

* Added LangChain integration

* Fixed issue created by git checkin process

* Added ':' to characters to remove from end of file path

* Tested initial migration to LangChain, removed comments and logging used for debugging

* Tested initial migration to LangChain, removed comments and logging used for debugging

* Converted camelCase to snake_case

* Turns out we need the exception handling

* Testing Hugging Face Integrations via LangChain

* Added LangChain loadable models

* Renames "qa" prompt to "clarify", since it's used in the "clarify" step, asking for clarification

* Fixed loading model yaml files

* Fixed streaming

* Added modeldir cli option

* Fixed typing

* Fixed interaction with token logging

* Fix spelling + dependency issues + typing

* Fix spelling + tests

* Removed unneeded logging which caused test to fail

* Cleaned up code

* Incorporated feedback

- deleted unnecessary functions & logger.info
- used LangChain ChatLLM instead of LLM to naturally communicate with gpt-4
- deleted loading model from yaml file, as LC doesn't offer this for ChatModels

* Update gpt_engineer/steps.py

Co-authored-by: Anton Osika <[email protected]>

* Incorporated feedback

- Fixed failing test
- Removed parsing complexity by using # type: ignore
- Replace every ocurence of ai.last_message_content with its content

* Fixed test

* Update gpt_engineer/steps.py

---------

Co-authored-by: H <[email protected]>
Co-authored-by: Anton Osika <[email protected]>

gpt_engineer/ai.py

@@ -1,13 +1,27 @@
 from __future__ import annotations
 
+import json
 import logging
 
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import List, Optional, Union
 
 import openai
 import tiktoken
 
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.chat_models import ChatOpenAI
+from langchain.chat_models.base import BaseChatModel
+from langchain.schema import (
+    AIMessage,
+    HumanMessage,
+    SystemMessage,
+    messages_from_dict,
+    messages_to_dict,
+)
+
+Message = Union[AIMessage, HumanMessage, SystemMessage]
+
 logger = logging.getLogger(__name__)
 
 
@@ -23,72 +37,69 @@ class TokenUsage:
 
 
 class AI:
-    def __init__(self, model="gpt-4", temperature=0.1):
+    def __init__(self, model_name="gpt-4", temperature=0.1):
         self.temperature = temperature
-        self.model = model
+        self.model_name = fallback_model(model_name)
+        self.llm = create_chat_model(self.model_name, temperature)
+        self.tokenizer = get_tokenizer(self.model_name)
 
         # initialize token usage log
         self.cumulative_prompt_tokens = 0
         self.cumulative_completion_tokens = 0
         self.cumulative_total_tokens = 0
         self.token_usage_log = []
 
-        try:
-            self.tokenizer = tiktoken.encoding_for_model(model)
-        except KeyError:
-            logger.debug(
-                f"Tiktoken encoder for model {model} not found. Using "
-                "cl100k_base encoder instead. The results may therefore be "
-                "inaccurate and should only be used as estimate."
-            )
-            self.tokenizer = tiktoken.get_encoding("cl100k_base")
-
-    def start(self, system, user, step_name):
-        messages = [
-            {"role": "system", "content": system},
-            {"role": "user", "content": user},
+    def start(self, system: str, user: str, step_name: str) -> List[Message]:
+        messages: List[Message] = [
+            SystemMessage(content=system),
+            HumanMessage(content=user),
         ]
-
         return self.next(messages, step_name=step_name)
 
-    def fsystem(self, msg):
-        return {"role": "system", "content": msg}
+    def fsystem(self, msg: str) -> SystemMessage:
+        return SystemMessage(content=msg)
 
-    def fuser(self, msg):
-        return {"role": "user", "content": msg}
+    def fuser(self, msg: str) -> HumanMessage:
+        return HumanMessage(content=msg)
 
-    def fassistant(self, msg):
-        return {"role": "assistant", "content": msg}
+    def fassistant(self, msg: str) -> AIMessage:
+        return AIMessage(content=msg)
 
-    def next(self, messages: List[Dict[str, str]], prompt=None, *, step_name=None):
+    def next(
+        self,
+        messages: List[Message],
+        prompt: Optional[str] = None,
+        *,
+        step_name: str,
+    ) -> List[Message]:
         if prompt:
-            messages += [{"role": "user", "content": prompt}]
+            messages.append(self.fuser(prompt))
 
         logger.debug(f"Creating a new chat completion: {messages}")
-        response = openai.ChatCompletion.create(
-            messages=messages,
-            stream=True,
-            model=self.model,
-            temperature=self.temperature,
-        )
 
-        chat = []
-        for chunk in response:
-            delta = chunk["choices"][0]["delta"]  # type: ignore
-            msg = delta.get("content", "")
-            print(msg, end="")
-            chat.append(msg)
-        print()
-        messages += [{"role": "assistant", "content": "".join(chat)}]
+        callsbacks = [StreamingStdOutCallbackHandler()]
+        response = self.llm(messages, callbacks=callsbacks)  # type: ignore
+        messages.append(response)
+
         logger.debug(f"Chat completion finished: {messages}")
 
         self.update_token_usage_log(
-            messages=messages, answer="".join(chat), step_name=step_name
+            messages=messages, answer=response.content, step_name=step_name
         )
 
         return messages
 
-    def update_token_usage_log(self, messages, answer, step_name):
+    @staticmethod
+    def serialize_messages(messages: List[Message]) -> str:
+        return json.dumps(messages_to_dict(messages))
+
+    @staticmethod
+    def deserialize_messages(jsondictstr: str) -> List[Message]:
+        return list(messages_from_dict(json.loads(jsondictstr)))  # type: ignore
+
+    def update_token_usage_log(
+        self, messages: List[Message], answer: str, step_name: str
+    ) -> None:
         prompt_tokens = self.num_tokens_from_messages(messages)
         completion_tokens = self.num_tokens(answer)
         total_tokens = prompt_tokens + completion_tokens
@@ -109,7 +120,7 @@ def update_token_usage_log(self, messages, answer, step_name):
             )
         )
 
-    def format_token_usage_log(self):
+    def format_token_usage_log(self) -> str:
         result = "step_name,"
         result += "prompt_tokens_in_step,completion_tokens_in_step,total_tokens_in_step"
         result += ",total_prompt_tokens,total_completion_tokens,total_tokens\n"
@@ -123,20 +134,17 @@ def format_token_usage_log(self):
             result += str(log.total_tokens) + "\n"
         return result
 
-    def num_tokens(self, txt):
+    def num_tokens(self, txt: str) -> int:
         return len(self.tokenizer.encode(txt))
 
-    def num_tokens_from_messages(self, messages):
+    def num_tokens_from_messages(self, messages: List[Message]) -> int:
         """Returns the number of tokens used by a list of messages."""
         n_tokens = 0
         for message in messages:
             n_tokens += (
                 4  # every message follows <im_start>{role/name}\n{content}<im_end>\n
             )
-            for key, value in message.items():
-                n_tokens += self.num_tokens(value)
-                if key == "name":  # if there's a name, the role is omitted
-                    n_tokens += -1  # role is always required and always 1 token
+            n_tokens += self.num_tokens(message.content)
         n_tokens += 2  # every reply is primed with <im_start>assistant
         return n_tokens
 
@@ -151,4 +159,39 @@ def fallback_model(model: str) -> str:
             "to gpt-3.5-turbo. Sign up for the GPT-4 wait list here: "
             "https://openai.com/waitlist/gpt-4-api\n"
         )
-        return "gpt-3.5-turbo-16k"
+        return "gpt-3.5-turbo"
+
+
+def create_chat_model(model: str, temperature) -> BaseChatModel:
+    if model == "gpt-4":
+        return ChatOpenAI(
+            model="gpt-4",
+            temperature=temperature,
+            streaming=True,
+            client=openai.ChatCompletion,
+        )
+    elif model == "gpt-3.5-turbo":
+        return ChatOpenAI(
+            model="gpt-3.5-turbo",
+            temperature=temperature,
+            streaming=True,
+            client=openai.ChatCompletion,
+        )
+    else:
+        raise ValueError(f"Model {model} is not supported.")
+
+
+def get_tokenizer(model: str):
+    if "gpt-4" in model or "gpt-3.5" in model:
+        return tiktoken.encoding_for_model(model)
+
+    logger.debug(
+        f"No encoder implemented for model {model}."
+        "Defaulting to tiktoken cl100k_base encoder."
+        "Use results only as estimates."
+    )
+    return tiktoken.get_encoding("cl100k_base")
+
+
+def serialize_messages(messages: List[Message]) -> str:
+    return AI.serialize_messages(messages)

gpt_engineer/chat_to_files.py

@@ -9,7 +9,7 @@ def parse_chat(chat):  # -> List[Tuple[str, str]]:
     files = []
     for match in matches:
         # Strip the filename of any non-allowed characters and convert / to \
-        path = re.sub(r'[<>"|?*]', "", match.group(1))
+        path = re.sub(r'[\:<>"|?*]', "", match.group(1))
 
         # Remove leading and trailing brackets
         path = re.sub(r"^\[(.*)\]$", r"\1", path)
@@ -18,7 +18,7 @@ def parse_chat(chat):  # -> List[Tuple[str, str]]:
         path = re.sub(r"^`(.*)`$", r"\1", path)
 
         # Remove trailing ]
-        path = re.sub(r"\]$", "", path)
+        path = re.sub(r"[\]\:]$", "", path)
 
         # Get the code
         code = match.group(2)

gpt_engineer/learning.py

@@ -100,11 +100,11 @@ def check_consent():
     path = Path(".gpte_consent")
     if path.exists() and path.read_text() == "true":
         return
-    ans = input("Is it ok if we store your prompts to learn? (y/n)")
-    while ans.lower() not in ("y", "n"):
-        ans = input("Invalid input. Please enter y or n: ")
+    answer = input("Is it ok if we store your prompts to learn? (y/n)")
+    while answer.lower() not in ("y", "n"):
+        answer = input("Invalid input. Please enter y or n: ")
 
-    if ans.lower() == "y":
+    if answer.lower() == "y":
         path.write_text("true")
         print(colored("Thank you️", "light_green"))
         print()
@@ -153,21 +153,14 @@ def ask_if_can_store() -> bool:
     return can_store == "y"
 
 
-def logs_to_string(steps: List[Step], logs: DB):
+def logs_to_string(steps: List[Step], logs: DB) -> str:
     chunks = []
     for step in steps:
         chunks.append(f"--- {step.__name__} ---\n")
-        messages = json.loads(logs[step.__name__])
-        chunks.append(format_messages(messages))
+        chunks.append(logs[step.__name__])
     return "\n".join(chunks)
 
 
-def format_messages(messages: List[dict]) -> str:
-    return "\n".join(
-        [f"{message['role']}:\n\n{message['content']}" for message in messages]
-    )
-
-
 def extract_learning(
     model: str, temperature: float, steps: List[Step], dbs: DBs, steps_file_hash
 ) -> Learning:

gpt_engineer/main.py

@@ -1,4 +1,3 @@
-import json
 import logging
 
 from pathlib import Path
@@ -28,7 +27,7 @@ def main(
 
     model = fallback_model(model)
     ai = AI(
-        model=model,
+        model_name=model,
         temperature=temperature,
     )
 
@@ -56,7 +55,7 @@ def main(
     steps = STEPS[steps_config]
     for step in steps:
         messages = step(ai, dbs)
-        dbs.logs[step.__name__] = json.dumps(messages)
+        dbs.logs[step.__name__] = AI.serialize_messages(messages)
 
     if collect_consent():
         collect_learnings(model, temperature, steps, dbs)