Try/catch in chatters

jrrobison1 · Apr 4, 2024 · 0461c4f · 0461c4f
1 parent 7a3852c
commit 0461c4f
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 17 deletions.
diff --git a/llm_writer_workshop/chatter/mistral_chatter.py b/llm_writer_workshop/chatter/mistral_chatter.py
@@ -18,15 +18,18 @@ def __init__(
         temperature=DEFAULT_TEMPERATURE,
         name
     ):
-        self.mistral_api_key = os.environ.get("MISTRAL_API_KEY")
-        self.model = model
-        self.max_output_tokens = max_output_tokens
-        self.temperature = temperature
-        self.name = name
-        self.message_history = []
-        self.client = MistralClient(api_key=self.mistral_api_key)
-        self.initial_prompt = ChatMessage(role="system", content=system_prompt)
-        self.message_history.append(self.initial_prompt)
+        try:
+            self.mistral_api_key = os.environ.get("MISTRAL_API_KEY")
+            self.model = model
+            self.max_output_tokens = max_output_tokens
+            self.temperature = temperature
+            self.name = name
+            self.message_history = []
+            self.client = MistralClient(api_key=self.mistral_api_key)
+            self.initial_prompt = ChatMessage(role="system", content=system_prompt)
+            self.message_history.append(self.initial_prompt)
+        except Exception as e:
+            print(e)
 
     def chat(self, message):
         if message != "":

diff --git a/llm_writer_workshop/chatter/openai_chatter.py b/llm_writer_workshop/chatter/openai_chatter.py
@@ -16,14 +16,17 @@ def __init__(
         name,
         max_output_tokens=DEFAULT_SYSTEM_PROMPT
     ):
-        self.model = model
-        self.max_output_tokens = max_output_tokens
-        self.name = name
-        self.initial_prompt = {"role": "system", "content": system_prompt}
-        self.message_history = []
-        self.message_history.append(self.initial_prompt)
-        self.openai_api_key = os.environ.get("OPENAI_API_KEY")
-        self.open_ai_client = OpenAI(api_key=self.openai_api_key)
+        try:
+            self.model = model
+            self.max_output_tokens = max_output_tokens
+            self.name = name
+            self.initial_prompt = {"role": "system", "content": system_prompt}
+            self.message_history = []
+            self.message_history.append(self.initial_prompt)
+            self.openai_api_key = os.environ.get("OPENAI_API_KEY")
+            self.open_ai_client = OpenAI(api_key=self.openai_api_key)
+        except Exception as e:
+            print(e)
 
     def chat(self, message):
         user_message = {}