import streamlit as st
import langchain
from langchain_community.chat_models import ChatOllama
from langchain.cache import InMemoryCache
from dotenv import load_dotenv
from langchain_community.embeddings import OllamaEmbeddings
import os
from PIL import Image
from chroma_main import answer_no_retriever

langchain.cache =  InMemoryCache()

load_dotenv()

CHROMA_DB = "./chroma_db"
MODEL = os.getenv("MODEL", "llama2")

OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
llm = ChatOllama(base_url=OLLAMA_BASE_URL, model=MODEL, temperature=0.0)
ollama_embeddings = OllamaEmbeddings(base_url=OLLAMA_BASE_URL, model="codellama")



st.button("clear history", type="primary")    
if st.button:
        st.session_state.messages = []


if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


if prompt := st.chat_input("What is your query?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").markdown(prompt)    
    
    response = answer_no_retriever(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        # message_placeholder.markdown(response + "▌")
        message_placeholder.markdown(response)
    st.session_state.messages.append({"role": "assistant", "content": response})
    

langchain_community.embeddings.OllamaEmbeddings langchain_community.chat_models.ChatOllama langchain.cache.InMemoryCache
[((324, 339), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (337, 339), False, 'from langchain.cache import InMemoryCache\n'), ((341, 354), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (352, 354), False, 'from dotenv import load_dotenv\n'), ((390, 418), 'os.getenv', 'os.getenv', (['"""MODEL"""', '"""llama2"""'], {}), "('MODEL', 'llama2')\n", (399, 418), False, 'import os\n'), ((438, 492), 'os.getenv', 'os.getenv', (['"""OLLAMA_BASE_URL"""', '"""http://localhost:11434"""'], {}), "('OLLAMA_BASE_URL', 'http://localhost:11434')\n", (447, 492), False, 'import os\n'), ((499, 565), 'langchain_community.chat_models.ChatOllama', 'ChatOllama', ([], {'base_url': 'OLLAMA_BASE_URL', 'model': 'MODEL', 'temperature': '(0.0)'}), '(base_url=OLLAMA_BASE_URL, model=MODEL, temperature=0.0)\n', (509, 565), False, 'from langchain_community.chat_models import ChatOllama\n'), ((586, 647), 'langchain_community.embeddings.OllamaEmbeddings', 'OllamaEmbeddings', ([], {'base_url': 'OLLAMA_BASE_URL', 'model': '"""codellama"""'}), "(base_url=OLLAMA_BASE_URL, model='codellama')\n", (602, 647), False, 'from langchain_community.embeddings import OllamaEmbeddings\n'), ((651, 693), 'streamlit.button', 'st.button', (['"""clear history"""'], {'type': '"""primary"""'}), "('clear history', type='primary')\n", (660, 693), True, 'import streamlit as st\n'), ((968, 1004), 'streamlit.chat_input', 'st.chat_input', (['"""What is your query?"""'], {}), "('What is your query?')\n", (981, 1004), True, 'import streamlit as st\n'), ((1010, 1079), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'user', 'content': prompt}"], {}), "({'role': 'user', 'content': prompt})\n", (1042, 1079), True, 'import streamlit as st\n'), ((1149, 1176), 'chroma_main.answer_no_retriever', 'answer_no_retriever', (['prompt'], {}), '(prompt)\n', (1168, 1176), False, 'from chroma_main import answer_no_retriever\n'), ((1364, 1440), 'streamlit.session_state.messages.append', 'st.session_state.messages.append', (["{'role': 'assistant', 'content': response}"], {}), "({'role': 'assistant', 'content': response})\n", (1396, 1440), True, 'import streamlit as st\n'), ((879, 911), 'streamlit.chat_message', 'st.chat_message', (["message['role']"], {}), "(message['role'])\n", (894, 911), True, 'import streamlit as st\n'), ((921, 952), 'streamlit.markdown', 'st.markdown', (["message['content']"], {}), "(message['content'])\n", (932, 952), True, 'import streamlit as st\n'), ((1187, 1215), 'streamlit.chat_message', 'st.chat_message', (['"""assistant"""'], {}), "('assistant')\n", (1202, 1215), True, 'import streamlit as st\n'), ((1247, 1257), 'streamlit.empty', 'st.empty', ([], {}), '()\n', (1255, 1257), True, 'import streamlit as st\n'), ((1084, 1107), 'streamlit.chat_message', 'st.chat_message', (['"""user"""'], {}), "('user')\n", (1099, 1107), True, 'import streamlit as st\n')]