feat: replace postgres with sqlite

app.py

@@ -2,6 +2,7 @@ import json
 import os
 
 import sqlalchemy
+import sqlite_vss
 import streamlit as st
 import streamlit.components.v1 as components
 from langchain import OpenAI
@@ -9,13 +10,14 @@ from langchain.callbacks import get_openai_callback
 from langchain.chains import ConversationalRetrievalChain
 from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.embeddings import GPT4AllEmbeddings
+from sqlalchemy import event
 
 from chat_history import insert_chat_history, insert_chat_history_articles
 from css import load_css
 from custom_pgvector import CustomPGVector
 from message import Message
 
-CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"
+CONNECTION_STRING = "sqlite:///data/sorbobot.db"
 
 st.set_page_config(layout="wide")
 
@@ -26,10 +28,16 @@ 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)
+
     conn = engine.connect()
     return conn
 
-
 conn = connect()
 
 

custom_pgvector.py

@@ -4,6 +4,7 @@ import contextlib
 import enum
 import json
 import logging
+import struct
 from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Type
 
 import pandas as pd
@@ -348,33 +349,54 @@ class CustomPGVector(VectorStore):
         k: int = 4,
     ) -> List[Any]:
         """Query the collection."""
-        with Session(self._conn) as session:
-            results = session.execute(
-                text(
-                    f"""
+        vector = bytearray(struct.pack("f" * len(embedding), *embedding))
+
+        cursor = self._conn.execute(
+            text("""
+                with matches as (
                     select
-                        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};
-                """
+                        rowid,
+                        distance
+                    from vss_article
+                    where vss_search(
+                        abstract_embedding,
+                        :vector
+                    )
+                    limit :limit
                 )
-            )
-            results = results.fetchall()
-            results = pd.DataFrame(results, columns=["id", "title", "doi", "abstract", "keywords",  "authors", "distance"])
-            results = results.to_dict(orient="records")
+                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.to_dict(orient="records")
         return results
 
     def similarity_search_by_vector(