Save messages on success

chain.py

@@ -7,6 +7,9 @@
 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 @@ def think(cls, cache: Conversation, input: str):
         ])
         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 @@ def revise_thought(cls, cache: Conversation, input: str, thought: str):
         ])
         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 @@ def respond(cls, cache: Conversation, thought: str, input: str):
         ])
         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 @@ async def __anext__(self):
         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