From c6d0aef2385a77d2ae6944f1a4bbacc07a80494d Mon Sep 17 00:00:00 2001
From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Date: Mon, 18 Dec 2023 16:14:17 -0500
Subject: [PATCH] Save messages on success

---
 chain.py | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/chain.py b/chain.py
index d4065849..b5885ae2 100644
--- a/chain.py
+++ b/chain.py
@@ -7,6 +7,9 @@ from dotenv import load_dotenv
 from collections.abc import AsyncIterator
 from cache import Conversation
 from typing import List
+
+from openai import BadRequestError
+
 import sentry_sdk
 
 load_dotenv()
@@ -52,11 +55,14 @@ class BloomChain:
         ])
         chain = thought_prompt | cls.llm 
 
-        cache.add_message("thought", HumanMessage(content=input))
 
+        def save_new_messages(ai_response):
+            cache.add_message("thought", HumanMessage(content=input))
+            cache.add_message("thought", AIMessage(content=ai_response))
+        
         return Streamable(
             chain.astream({}, {"tags": ["thought"], "metadata": {"conversation_id": cache.conversation_id, "user_id": cache.user_id}}),
-            lambda thought: cache.add_message("thought", AIMessage(content=thought)) 
+            save_new_messages 
         )
     
     @classmethod
@@ -75,11 +81,13 @@ class BloomChain:
         ])
         chain = messages | cls.llm
 
-        cache.add_message("thought_revision", HumanMessage(content=input))
+        def save_new_messages(ai_response):
+            cache.add_message("thought_revision", HumanMessage(content=input))
+            cache.add_message("thought_revision", AIMessage(content=ai_response))
 
         return Streamable(
             chain.astream({ "thought": thought, "retrieved_vectors": "\n".join(doc.page_content for doc in docs)}, {"tags": ["thought_revision"], "metadata": {"conversation_id": cache.conversation_id, "user_id": cache.user_id}}),
-            lambda thought_revision: cache.add_message("thought_revision", AIMessage(content=thought_revision))  # add the revised thought to thought memory
+            save_new_messages
         )
         
     @classmethod
@@ -93,11 +101,13 @@ class BloomChain:
         ])
         chain = response_prompt | cls.llm
 
-        cache.add_message("response", HumanMessage(content=input))
+        def save_new_messages(ai_response):
+            cache.add_message("response", HumanMessage(content=input))
+            cache.add_message("response", AIMessage(content=ai_response))
 
         return Streamable(
             chain.astream({ "thought": thought }, {"tags": ["response"], "metadata": {"conversation_id": cache.conversation_id, "user_id": cache.user_id}}),
-            lambda response: cache.add_message("response", AIMessage(content=response))
+            save_new_messages
         )
     
     @classmethod
@@ -337,6 +347,12 @@ class Streamable:
         except StopAsyncIteration as e:
             self.callback(self.content)
             raise StopAsyncIteration
+        except BadRequestError as e:
+            if e.code == "content_filter":
+                self.stream_error = True
+                self.message = "Sorry, your message was flagged as inappropriate. Please try again."
+
+                return self.message
         except Exception as e:
             raise e