feat: move sql connection outside of custom pgvector

app.py

@@ -11,6 +11,7 @@ from langchain.chains import ConversationalRetrievalChain
 from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.embeddings import GPT4AllEmbeddings
 from message import Message
+import sqlalchemy
 
 CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"
 
@@ -21,6 +22,15 @@ st.title("Sorbobot - Le futur de la recherche scientifique interactive")
 chat_column, doc_column = st.columns([2, 1])
 
 
+def connect() -> sqlalchemy.engine.Connection:
+    engine = sqlalchemy.create_engine(CONNECTION_STRING)
+    conn = engine.connect()
+    return conn
+
+
+conn = connect()
+
+
 def initialize_session_state():
     if "history" not in st.session_state:
         st.session_state.history = []
@@ -33,7 +43,7 @@ def initialize_session_state():
             embedding_function=embeddings,
             table_name="article",
             column_name="abstract_embedding",
-            connection_string=CONNECTION_STRING,
+            connection=conn,
         )
 
         retriever = db.as_retriever()

custom_pgvector.py

@@ -70,7 +70,7 @@ class CustomPGVector(VectorStore):
     To use, you should have the ``pgvector`` python package installed.
 
     Args:
-        connection_string: Postgres connection string.
+        connection: Postgres connection string.
         embedding_function: Any embedding function implementing
             `langchain.embeddings.base.Embeddings` interface.
         table_name: The name of the collection to use. (default: langchain)
@@ -87,14 +87,13 @@ class CustomPGVector(VectorStore):
             from langchain.vectorstores import PGVector
             from langchain.embeddings.openai import OpenAIEmbeddings
 
-            CONNECTION_STRING = "postgresql+psycopg2://hwc@localhost:5432/test3"
             COLLECTION_NAME = "state_of_the_union_test"
             embeddings = OpenAIEmbeddings()
             vectorestore = PGVector.from_documents(
                 embedding=embeddings,
                 documents=docs,
                 table_name=COLLECTION_NAME,
-                connection_string=CONNECTION_STRING,
+                connection=connection,
             )
 
 
@@ -102,7 +101,7 @@ class CustomPGVector(VectorStore):
 
     def __init__(
         self,
-        connection_string: str,
+        connection: sqlalchemy.engine.Connection,
         embedding_function: Embeddings,
         table_name: str,
         column_name: str,
@@ -111,7 +110,7 @@ class CustomPGVector(VectorStore):
         pre_delete_collection: bool = False,
         logger: Optional[logging.Logger] = None,
     ) -> None:
-        self.connection_string = connection_string
+        self._conn = connection
         self.embedding_function = embedding_function
         self.table_name = table_name
         self.column_name = column_name
@@ -127,7 +126,7 @@ class CustomPGVector(VectorStore):
         """
         Initialize the store.
         """
-        self._conn = self.connect()
+        # self._conn = self.connect()
         self.create_vector_extension()
 
         self.EmbeddingStore = Article
@@ -136,11 +135,6 @@ class CustomPGVector(VectorStore):
     def embeddings(self) -> Embeddings:
         return self.embedding_function
 
-    def connect(self) -> sqlalchemy.engine.Connection:
-        engine = sqlalchemy.create_engine(self.connection_string)
-        conn = engine.connect()
-        return conn
-
     def create_vector_extension(self) -> None:
         try:
             with Session(self._conn) as session: