feat: Add sqlite-vss for vector search similariy capabilities

* feat: Add sqlite-vss to add simioary search to sqlite

* feat: Integrate with langchain for sqlite-vss implementation

* feat: VSS now fully working

* fix: remove unsused file

* fix: Adjust unit test for similarity search

Author: philippe2803

.gitignore

@@ -1,4 +1,5 @@
 contentmap.db
+.DS_Store
 /scratch
 
  Byte-compiled / optimized / DLL files

Dockerfile

@@ -18,4 +18,4 @@ RUN python3 -c 'from sentence_transformers import SentenceTransformer; embedder
 
 ADD . /app
 
-CMD ["pytest", "./tests"]
+CMD ["pytest", "./tests"]

contentmap/__init__.py

@@ -0,0 +1,2 @@
+from contentmap.core import ContentMapCreator
+from contentmap.sitemap import SitemapToContentDatabase

contentmap/core.py

@@ -1,37 +1,44 @@
-from typing import List, Dict, TypedDict
+from typing import List, Dict
 from datetime import datetime
 import sqlite3
-import importlib.metadata
+from contentmap.vss import ContentMapVSS
 
 
 class ContentMapCreator:
 
     def __init__(
             self,
             contents: List[Dict[str, str]],
-            database: str = "contentmap.db"
+            database: str = "contentmap.db",
+            include_vss: bool = False
     ):
         self.contents = contents
-        self.db = sqlite3.connect(database)
-        self.db.row_factory = sqlite3.Row
-        self.cursor = self.db.cursor()
+        self.include_vss = include_vss
+        self.connection = sqlite3.connect(database)
+        self.connection.row_factory = sqlite3.Row
+
+        if self.include_vss:
+            import sqlite_vss
+            self.connection.enable_load_extension(True)
+            sqlite_vss.load(self.connection)
+            self.connection.enable_load_extension(False)
+
+        self.cursor = self.connection.cursor()
 
     def init_db(self):
         self.cursor.execute("CREATE TABLE IF NOT EXISTS content (url, content)")
         self.cursor.execute("CREATE TABLE IF NOT EXISTS config (cat, value)")
-        self.db.commit()
+        self.connection.commit()
 
     def add_config(self):
         data = [
-            {"Version:": "1"},
             {"Generated with:": "Contentmap lib"},
             {"Date:": datetime.now().strftime("%Y-%m-%d %H:%M:%S")},
-            {"Embeddings:": "mistral-embed"},
-            {"FTSE:": None},
+            {"Embeddings:": "all-MiniLM-L6-v2"}
         ]
         data = [{"cat": k, "value": v} for row in data for k, v in row.items()]
         self.cursor.executemany("INSERT INTO config VALUES (:cat, :value)", data)
-        self.db.commit()
+        self.connection.commit()
 
     def build(self):
         self.init_db()
@@ -40,4 +47,11 @@ def build(self):
             "INSERT INTO content VALUES (:url, :content)",
             self.contents
         )
-        self.db.commit()
+        self.connection.commit()
+
+        if self.include_vss:
+            self.add_vss()
+
+    def add_vss(self):
+        cm_vss = ContentMapVSS(connection=self.connection)
+        cm_vss.load()

contentmap/ftse.py

@@ -12,19 +12,21 @@
 
 class SitemapToContentDatabase:
 
-    def __init__(self, sitemap_url, seconds_timeout=10, concurrency=None):
+    def __init__(self, sitemap_url, seconds_timeout=10, concurrency=None,
+                 include_vss=False):
         self.sitemap_url = sitemap_url
         self.semaphore = asyncio.Semaphore(concurrency) if concurrency is not None else None
         self.timeout = aiohttp.ClientTimeout(
             sock_connect=seconds_timeout,
             sock_read=seconds_timeout
         )
+        self.include_vss = include_vss
 
-    def load(self):
+    def build(self):
         urls = self.get_urls()
         loop = asyncio.get_event_loop()
         contents = loop.run_until_complete(self.get_contents(urls))
-        cm = ContentMapCreator(contents)
+        cm = ContentMapCreator(contents, include_vss=self.include_vss)
         cm.build()
 
     def get_urls(self):

contentmap/sitemap.py

@@ -5,10 +5,7 @@
 import sqlite3
 from typing import Optional
 
-import sqlite_vss
-from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.document_loaders import TextLoader
 from langchain_community.embeddings.sentence_transformer import (
     SentenceTransformerEmbeddings,
 )
@@ -59,8 +56,10 @@ def prepare_texts_and_metadatas(self):
         rows = result.fetchall()
 
         # based on Anyscale analysis (https://t.ly/yjgxQ), it looks like the
-        # sweet spot is 700 chunk size and 50 chunk overlap
-        text_splitter = CharacterTextSplitter(chunk_size=700, chunk_overlap=50)
+        # sweet spot is 700 chunk size and 50 chunk overlap.
+        text_splitter = CharacterTextSplitter(
+            chunk_size=700, chunk_overlap=50, separator="."
+        )
 
         texts = []
         metadatas = []
@@ -73,4 +72,9 @@ def prepare_texts_and_metadatas(self):
         return texts, metadatas
 
     def similarity_search(self, *args, **kwargs):
-        return self.vss.similarity_search(*args, **kwargs)
+        data = self.vss.similarity_search(*args, **kwargs)
+        rag_results = []
+        for doc in data:
+            item = {"content": doc.page_content, "url": doc.metadata['url']}
+            rag_results.append(item)
+        return rag_results

contentmap/vss.py

@@ -0,0 +1,6 @@
+version: "3"
+services:
+  app:
+    build:
+      context: .
+    image: contentmap_local