Save messages on success

This commit is contained in:
Vineeth Voruganti 2023-12-18 16:14:17 -05:00
parent cda27b5bc2
commit c6d0aef238
1 changed files with 22 additions and 6 deletions

View File

@ -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