feat: first draft of conversational retrieval

app.py

@@ -5,10 +5,15 @@ import streamlit.components.v1 as components
 from css import load_css
 from langchain import OpenAI
 from langchain.callbacks import get_openai_callback
-from langchain.chains import ConversationChain
-from langchain.chains.conversation.memory import ConversationSummaryMemory
+from langchain.chains import ConversationalRetrievalChain
+from langchain.chains.conversation.memory import ConversationBufferMemory
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores.pgvector import PGVector
 from message import Message
 
+CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"
+COLLECTION_NAME = ""
+
 
 def initialize_session_state():
     if "history" not in st.session_state:
@@ -16,21 +21,37 @@ def initialize_session_state():
     if "token_count" not in st.session_state:
         st.session_state.token_count = 0
     if "conversation" not in st.session_state:
+        embeddings = OpenAIEmbeddings()
+
+        store = PGVector(
+            collection_name=COLLECTION_NAME,
+            connection_string=CONNECTION_STRING,
+            embedding_function=embeddings,
+        )
+
+        retriever = store.as_retriever()
+
         llm = OpenAI(
             temperature=0,
             openai_api_key=os.environ["OPENAI_API_KEY"],
-            model_name="text-davinci-003",
+            model="text-davinci-003",
         )
-        st.session_state.conversation = ConversationChain(
-            llm=llm,
-            memory=ConversationSummaryMemory(llm=llm),
+
+        st.session_state.memory = ConversationBufferMemory()
+        st.session_state.conversation = ConversationalRetrievalChain.from_llm(
+            llm=llm, retriever=retriever
         )
 
 
 def on_click_callback():
     with get_openai_callback() as cb:
         human_prompt = st.session_state.human_prompt
-        llm_response = st.session_state.conversation.run(human_prompt)
+        llm_response = st.session_state.conversation.run(
+            {
+                "question": human_prompt,
+                "chat_history": st.session_state.memory.buffer,
+            }
+        )
         st.session_state.history.append(Message("human", human_prompt))
         st.session_state.history.append(Message("ai", llm_response))
         st.session_state.token_count += cb.total_tokens
@@ -84,7 +105,7 @@ information_placeholder.caption(
     f"""
 Used {st.session_state.token_count} tokens \n
 Debug Langchain conversation: 
-{st.session_state.conversation.memory.buffer}
+{st.session_state.memory.buffer}
 """
 )