feat: Add sqlite-vss to add similarity search to sqlite

.github/workflows/ci.yaml

@@ -13,7 +13,7 @@ jobs:
     strategy:
         fail-fast: true
         matrix:
-          python-version: ["3.8", "3.9", "3.10"]
+          python-version: ["3.8", "3.9", "3.10", "3.11"]
 
     steps:
     - uses: actions/checkout@v3

.gitignore

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

Dockerfile

@@ -0,0 +1,21 @@
+FROM ubuntu:22.04
+
+
+RUN apt update
+RUN apt install -y python3-pip libgomp1 libatlas-base-dev liblapack-dev libsqlite3-dev
+
+WORKDIR /app
+
+ADD poetry.lock /app/poetry.lock
+ADD pyproject.toml /app/pyproject.toml
+
+RUN pip install poetry
+RUN poetry config virtualenvs.create false
+RUN poetry install
+
+RUN python3 -c 'from sentence_transformers import SentenceTransformer; embedder = SentenceTransformer("all-MiniLM-L6-v2")'
+
+
+ADD . /app
+
+CMD ["pytest", "./tests"]

README.md

@@ -1,3 +1,29 @@
 # Content map
 
-A way to share content from a specific domain. Similar to sitemaps but for content.
+A way to share content from a specific domain using SQLite as an alternative to 
+RSS feeds. The purpose of this library is to simply create a dataset for all the
+content on your website, using the XML sitemap as a starting point.
+
+
+## Installation
+
+```bash
+
+pip install contentmap
+
+```
+
+## Quickstart
+
+To build your contentmap.db that will contain all your content using your XML 
+sitemap as a starting point, you only need to write the following: 
+
+```python
+from contentmap.sitemap import SitemapToContentDatabase
+
+database = SitemapToContentDatabase("https://yourblog.com/sitemap.xml")
+database.load()
+
+```
+
+You can control how many urls can be crawled concurrently and also set some timeout.

contentmap/vss.py

@@ -0,0 +1,76 @@
+"""
+Class ContentMapVSS to create vector search dataset from a contentmap
+dataset already created.
+"""
+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,
+)
+from langchain_community.vectorstores import SQLiteVSS
+
+
+class ContentMapVSS:
+
+    def __init__(self,
+                 connection: Optional[sqlite3.Connection] = None,
+                 db_file: str = "contentmap.db"
+                 ):
+
+        self.connection = connection
+        if not connection:
+            self.connection = SQLiteVSS.create_connection(db_file)
+
+        embedding_function = SentenceTransformerEmbeddings(
+            model_name="all-MiniLM-L6-v2"
+        )
+        self.vss = SQLiteVSS(
+            table="content_chunks",
+            embedding=embedding_function,
+            connection=self.connection
+        )
+
+    def load(self):
+        # content table must be there
+        assert self.table_exists(table_name="content")
+        texts, metadatas = self.prepare_texts_and_metadatas()
+        self.vss.add_texts(texts=texts, metadatas=metadatas)
+        return self.vss
+
+    def table_exists(self, table_name: str) -> bool:
+        res = self.connection.execute(f"""
+            SELECT name 
+            FROM sqlite_master 
+            WHERE type='table' AND name='{table_name}';
+        """)
+        rows = res.fetchall()
+        if len(rows) == 1:
+            return True
+        return False
+
+    def prepare_texts_and_metadatas(self):
+        cursor = self.connection.cursor()
+        result = cursor.execute("SELECT content, url FROM content")
+        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)
+
+        texts = []
+        metadatas = []
+        for row in rows:
+            chunks = text_splitter.split_text(row["content"])
+            chunk_metadatas = [{"url": row["url"]} for _ in chunks]
+            texts += chunks
+            metadatas += chunk_metadatas
+
+        return texts, metadatas
+
+    def similarity_search(self, *args, **kwargs):
+        return self.vss.similarity_search(*args, **kwargs)

pyproject.toml

@@ -14,6 +14,9 @@ tqdm = "^4.66.1"
 lxml = "4.9.4"
 trafilatura = "^1.6.4"
 aiohttp = "^3.9.1"
+sqlite-vss = "^0.1.2"
+langchain = "^0.1.8"
+sentence-transformers = "^2.3.1"
 
 
 [tool.poetry.group.test.dependencies]

tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 import os
+import os.path as op
+import logging
 
 
 @pytest.fixture(autouse=True)
@@ -8,5 +10,9 @@ def remove_created_database_after_test():
     # Setup logic
     yield   # this is where the testing happens
     # Teardown logic
-    os.remove("contentmap.db")
+
+    contentmap_db_path = op.join(op.dirname(__file__), "contentmap.db")
+    if op.exists(contentmap_db_path):
+        logging.info('Destroying mock sqlite content instance')
+        os.remove(contentmap_db_path)
 

tests/test_fts.py

@@ -0,0 +1,12 @@
+import sqlite3
+
+
+def test_fts_extension_enabled():
+
+    con = sqlite3.connect(':memory:')
+    cur = con.cursor()
+    cur.execute('pragma compile_options;')
+    available_pragmas = cur.fetchall()
+    con.close()
+
+    assert ('ENABLE_FTS5',) in available_pragmas

tests/test_vss.py

@@ -0,0 +1,50 @@
+from contentmap.vss import ContentMapVSS
+import os.path as op
+from tests.utils import build_fixture_db
+
+
+class TestContentMapVSS:
+
+    def test_assertion_content_exists(self):
+        fixture_db = op.join(op.dirname(__file__), "fixture.db")
+        vss_content = ContentMapVSS(db_file=fixture_db)
+        assert vss_content.table_exists(table_name="content") is True
+
+    def test_assertion_content_not_exists(self):
+        vss_content = ContentMapVSS(db_file=":memory:")
+        assert vss_content.table_exists(table_name="content") is False
+
+
+class TestVssTablesCreation:
+
+    def test_vss_instance(self):
+        db = build_fixture_db()
+        cm_vss = ContentMapVSS(db_file=db)
+        cm_vss.load()
+        assert cm_vss.table_exists("content_chunks")
+
+    def test_prepare_texts_and_metadatas(self):
+        db = build_fixture_db()
+        cm_vss = ContentMapVSS(db_file=db)
+        texts, metadatas = cm_vss.prepare_texts_and_metadatas()
+        assert len(texts) == len(metadatas) >= 1
+
+    def test_chunk_table(self):
+        db = build_fixture_db()
+        cm_vss = ContentMapVSS(db_file=db)
+        cm_vss.load()
+        assert cm_vss.table_exists("content_chunks")
+        cursor = cm_vss.connection.cursor()
+        res = cursor.execute("SELECT * FROM content_chunks")
+        rows = res.fetchall()
+        assert len(rows) >= 15
+
+    def test_similarity_search(self):
+        db = build_fixture_db()
+        cm_vss = ContentMapVSS(db_file=db)
+        cm_vss.load()
+        data = cm_vss.similarity_search(query="who is Mistral ai company?", k=2)
+        assert len(data) == 2
+        metadatas = [doc.metadata for doc in data]
+        for meta in metadatas:
+            assert meta.get("url") == "https://philippeoger.com/pages/ai-scene-in-europe-last-week/"

tests/utils.py

@@ -0,0 +1,9 @@
+import os.path as op
+import shutil
+
+
+def build_fixture_db():
+    fixture_db = op.join(op.dirname(__file__), 'fixture.db')
+    dest = op.join(op.dirname(__file__), 'contentmap.db')
+    shutil.copy(fixture_db, dest)
+    return dest