Spaces:
Running
Running
| import streamlit as st | |
| from streamlit.logger import get_logger | |
| from langchain.memory import ConversationBufferMemory | |
| from models.databricks.instruction_memory import CustomBufferInstructionMemory | |
| from utils.mongo_utils import new_convo | |
| logger = get_logger(__name__) | |
| def push_convo2db(memories, username, language): | |
| if len(memories) == 1: | |
| issue = memories['memory']['issue'] | |
| model_one = memories['memory']['source'] | |
| new_convo(st.session_state['db_client'], issue, language, username, False, model_one) | |
| else: | |
| issue = memories['commonMemory']['issue'] | |
| model_one = memories['memoryA']['source'] | |
| model_two = memories['memoryB']['source'] | |
| new_convo(st.session_state['db_client'], issue, language, username, True, model_one, model_two) | |
| def change_memories(memories, language, changed_source=False): | |
| for memory, params in memories.items(): | |
| if (memory not in st.session_state) or changed_source: | |
| source = params['source'] | |
| logger.info(f"Source for memory {memory} is {source}") | |
| if source in ('OA_rolemodel','OA_finetuned',"CTL_llama2"): | |
| st.session_state[memory] = ConversationBufferMemory(ai_prefix='texter', human_prefix='helper') | |
| elif source in ('CTL_mistral'): | |
| st.session_state[memory] = CustomBufferInstructionMemory(human_prefix="</s> [INST]", memory_key="history") | |
| if ("convo_id" in st.session_state) and changed_source: | |
| del st.session_state['convo_id'] | |
| def clear_memory(memories, username, language): | |
| for memory, _ in memories.items(): | |
| st.session_state[memory].clear() | |
| if "convo_id" in st.session_state: | |
| del st.session_state['convo_id'] |