feat: postgres

app.py

@@ -17,7 +17,7 @@ from css import load_css
 from custom_pgvector import CustomPGVector
 from message import Message
 
-CONNECTION_STRING = "sqlite:///data/sorbobot.db"
+CONNECTION_STRING = "postgresql+psycopg2://postgres@/sorbobot?host=localhost"
 
 st.set_page_config(layout="wide")
 
@@ -27,13 +27,7 @@ chat_column, doc_column = st.columns([2, 1])
 
 
 def connect() -> sqlalchemy.engine.Connection:
-    engine = sqlalchemy.create_engine(CONNECTION_STRING)
-
-    @event.listens_for(engine, "connect")
-    def receive_connect(connection, _):
-        connection.enable_load_extension(True)
-        sqlite_vss.load(connection)
-        connection.enable_load_extension(False)
+    engine = sqlalchemy.create_engine(CONNECTION_STRING, pool_pre_ping=True)
 
     conn = engine.connect()
     return conn

custom_pgvector.py

@@ -339,53 +339,43 @@ class CustomPGVector(VectorStore):
         k: int = 4,
     ) -> List[Any]:
         """Query the collection."""
-        vector = bytearray(struct.pack("f" * len(embedding), *embedding))
-
-        cursor = self._conn.execute(
-            text("""
-                with matches as (
+        with Session(self._conn) as session:
+            results = session.execute(
+                text(
+                    f"""
                     select
-                        rowid,
-                        distance
-                    from vss_article
-                    where vss_search(
-                        abstract_embedding,
-                        :vector
-                    )
-                    limit 5
+                        a.id,
+                        a.title,
+                        a.doi,
+                        a.abstract,
+                        string_agg(distinct keyword."name", ',') as keywords, 
+                        string_agg(distinct author."name", ',') as authors,
+                        abstract_embedding <-> '{str(embedding)}' as distance
+                    from article a
+                    left join article_keyword ON article_keyword.article_id = a.id 
+                    left join keyword on article_keyword.keyword_id = keyword.id
+                    left join article_author ON article_author.article_id = a.id
+                    left join author on author.id = article_author.author_id  
+                    where abstract != 'NaN'
+                    GROUP BY a.id
+                    ORDER BY distance
+                    LIMIT {k};
+                """
                 )
-                select
-                    article.id,
-                    article.title,
-                    article.doi,
-                    article.abstract,
-                    group_concat(keyword."name", ',') as keywords,
-                    group_concat(author."name", ',') as authors,
-                    matches.distance
-                from matches
-                left join article on matches.rowid = article.rowid
-                left join article_keyword ak ON ak.article_id = article.id
-                left join keyword on ak.keyword_id = keyword.id
-                left join article_author ON article_author.article_id = article.id
-                left join author on author.id = article_author.author_id
-                group by article.id
-                order by distance;
-            """),
-            {"vector": vector, "limit": k}
-        )
-        results = cursor.fetchall()
-        results = pd.DataFrame(
-            results,
-            columns=[
-                "id",
-                "title",
-                "doi",
-                "abstract",
-                "keywords",
-                "authors",
-                "distance",
-            ],
-        )
+            )
+            results = results.fetchall()
+            results = pd.DataFrame(
+                results,
+                columns=[
+                    "id",
+                    "title",
+                    "doi",
+                    "abstract",
+                    "keywords",
+                    "authors",
+                    "distance",
+                ],
+            )
         results = results.to_dict(orient="records")
         return results