feat: move chain to specific file

sorbobotapp/app.py

@@ -5,17 +5,11 @@ import streamlit as st
 import streamlit.components.v1 as components
 from langchain.callbacks import get_openai_callback
 
-from langchain.chains.conversation.memory import ConversationBufferMemory
-from langchain.embeddings import GPT4AllEmbeddings
-from langchain.llms import OpenAI
-
+from chain import get_chain
 from chat_history import insert_chat_history, insert_chat_history_articles
 from connection import connect
 from css import load_css
 from message import Message
-from vector_store import CustomVectorStore
-from conversation_retrieval_chain import CustomConversationalRetrievalChain
-
 
 st.set_page_config(layout="wide")
 
@@ -32,34 +26,7 @@ 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 = GPT4AllEmbeddings()
-
-        db = CustomVectorStore(
-            embedding_function=embeddings,
-            table_name="article",
-            column_name="abstract_embedding",
-            connection=conn,
-        )
-
-        retriever = db.as_retriever()
-
-        llm = OpenAI(
-            temperature=0,
-            openai_api_key=os.environ["OPENAI_API_KEY"],
-            model="text-davinci-003",
-        )
-
-        memory = ConversationBufferMemory(
-            output_key="answer", memory_key="chat_history", return_messages=True
-        )
-        st.session_state.conversation = CustomConversationalRetrievalChain.from_llm(
-            llm=llm,
-            retriever=retriever,
-            verbose=True,
-            memory=memory,
-            return_source_documents=True,
-            max_tokens_limit=3700,
-        )
+        st.session_state.conversation = get_chain(conn)
 
 
 def send_message_callback():

sorbobotapp/chain.py

@@ -0,0 +1,40 @@
+import os
+
+import sqlalchemy
+from langchain.chains.conversation.memory import ConversationBufferMemory
+from langchain.embeddings import GPT4AllEmbeddings
+from langchain.llms import OpenAI
+
+from conversation_retrieval_chain import CustomConversationalRetrievalChain
+from vector_store import CustomVectorStore
+
+
+def get_chain(conn: sqlalchemy.engine.Connection):
+    embeddings = GPT4AllEmbeddings()
+
+    db = CustomVectorStore(
+        embedding_function=embeddings,
+        table_name="article",
+        column_name="abstract_embedding",
+        connection=conn,
+    )
+
+    retriever = db.as_retriever()
+
+    llm = OpenAI(
+        temperature=0,
+        openai_api_key=os.environ["OPENAI_API_KEY"],
+        model="text-davinci-003",
+    )
+
+    memory = ConversationBufferMemory(
+        output_key="answer", memory_key="chat_history", return_messages=True
+    )
+    return CustomConversationalRetrievalChain.from_llm(
+        llm=llm,
+        retriever=retriever,
+        verbose=True,
+        memory=memory,
+        return_source_documents=True,
+        max_tokens_limit=3700,
+    )