From c29cbd6858235cd8622eac6b0485c27c65b6f006 Mon Sep 17 00:00:00 2001 From: Manolo Santos Date: Wed, 11 Dec 2024 10:39:56 +0100 Subject: [PATCH 1/5] fix: Conditional parameters in pgvector. --- src/raglite/_database.py | 61 +++++++++++++++++++++++++++++----------- 1 file changed, 45 insertions(+), 16 deletions(-) diff --git a/src/raglite/_database.py b/src/raglite/_database.py index 573a3cc..9016c68 100644 --- a/src/raglite/_database.py +++ b/src/raglite/_database.py @@ -11,6 +11,7 @@ import numpy as np from markdown_it import MarkdownIt +from packaging import version from pydantic import ConfigDict from sqlalchemy.engine import Engine, make_url from sqlmodel import JSON, Column, Field, Relationship, Session, SQLModel, create_engine, text @@ -310,6 +311,18 @@ def from_chunks( ) +@lru_cache(maxsize=1) +def _get_pgvector_version(session: Session) -> str | None: + """Get pgvector version. + + Returns + ------- + str | None: Version string if pgvector is installed, None otherwise + """ + result = session.execute(text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")) + return result.scalar() + + @lru_cache(maxsize=1) def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: """Create a database engine and initialize it.""" @@ -354,21 +367,29 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: with Session(engine) as session: metrics = {"cosine": "cosine", "dot": "ip", "euclidean": "l2", "l1": "l1", "l2": "l2"} session.execute( - text(""" - CREATE INDEX IF NOT EXISTS keyword_search_chunk_index ON chunk USING GIN (to_tsvector('simple', body)); - """) + text( + """ + CREATE INDEX IF NOT EXISTS keyword_search_chunk_index ON chunk USING GIN (to_tsvector('simple', body)); + """ + ) ) - session.execute( - text(f""" + base_sql = f""" CREATE INDEX IF NOT EXISTS vector_search_chunk_index ON chunk_embedding USING hnsw ( (embedding::halfvec({embedding_dim})) halfvec_{metrics[config.vector_search_index_metric]}_ops ); SET hnsw.ef_search = {20 * 4 * 8}; - SET hnsw.iterative_scan = {'relaxed_order' if config.reranker else 'strict_order'}; - """) - ) + """ + # Add iterative scan if version >= 0.8.0 + pgvector_version = _get_pgvector_version(session) + if pgvector_version and version.parse(pgvector_version) >= version.parse("0.8.0"): + sql = f"""{base_sql}; + SET hnsw.iterative_scan = {'relaxed_order' if config.reranker else 'strict_order'}; + """ + else: + sql = f"{base_sql};" + session.execute(text(sql)) session.commit() elif db_backend == "sqlite": # Create a virtual table for keyword search on the chunk table. @@ -376,31 +397,39 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: # [1] https://www.sqlite.org/fts5.html#external_content_tables with Session(engine) as session: session.execute( - text(""" + text( + """ CREATE VIRTUAL TABLE IF NOT EXISTS keyword_search_chunk_index USING fts5(body, content='chunk', content_rowid='rowid'); - """) + """ + ) ) session.execute( - text(""" + text( + """ CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_insert AFTER INSERT ON chunk BEGIN INSERT INTO keyword_search_chunk_index(rowid, body) VALUES (new.rowid, new.body); END; - """) + """ + ) ) session.execute( - text(""" + text( + """ CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_delete AFTER DELETE ON chunk BEGIN INSERT INTO keyword_search_chunk_index(keyword_search_chunk_index, rowid, body) VALUES('delete', old.rowid, old.body); END; - """) + """ + ) ) session.execute( - text(""" + text( + """ CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_update AFTER UPDATE ON chunk BEGIN INSERT INTO keyword_search_chunk_index(keyword_search_chunk_index, rowid, body) VALUES('delete', old.rowid, old.body); INSERT INTO keyword_search_chunk_index(rowid, body) VALUES (new.rowid, new.body); END; - """) + """ + ) ) session.commit() return engine From 0b4220b1b4c01491800074ab94aa347299cce33e Mon Sep 17 00:00:00 2001 From: Manolo Santos Date: Wed, 11 Dec 2024 12:18:37 +0100 Subject: [PATCH 2/5] Update src/raglite/_database.py Co-authored-by: Laurent Sorber --- src/raglite/_database.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/raglite/_database.py b/src/raglite/_database.py index 9016c68..ecef10e 100644 --- a/src/raglite/_database.py +++ b/src/raglite/_database.py @@ -312,15 +312,14 @@ def from_chunks( @lru_cache(maxsize=1) -def _get_pgvector_version(session: Session) -> str | None: - """Get pgvector version. - - Returns - ------- - str | None: Version string if pgvector is installed, None otherwise - """ - result = session.execute(text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")) - return result.scalar() +def _pgvector_version(session: Session) -> Version: + try: + result = session.execute(text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")) + pgvector_version = version.parse(result.scalar()) + except Exception as e: + error_message = "Unable to parse pgvector version, is pgvector installed?" + raise ValueError(error_message) from e + return pgvector_version @lru_cache(maxsize=1) From 1e0e20ff61361650b6569887524875ce3dbb0ee8 Mon Sep 17 00:00:00 2001 From: Manolo Santos Date: Wed, 11 Dec 2024 14:33:18 +0100 Subject: [PATCH 3/5] chore: code style --- poetry.lock | 2 +- pyproject.toml | 2 ++ src/raglite/_database.py | 47 ++++++++++++++++------------------------ 3 files changed, 22 insertions(+), 29 deletions(-) diff --git a/poetry.lock b/poetry.lock index 7353b43..b207e92 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. [[package]] name = "aiofiles" diff --git a/pyproject.toml b/pyproject.toml index 10e3f57..fe253bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,8 @@ ragas = { version = ">=0.1.12", optional = true } typer = ">=0.12.5" # Frontend: chainlit = { version = ">=1.2.0", optional = true } +# Utilities: +packaging = ">=23.0" [tool.poetry.extras] # https://python-poetry.org/docs/pyproject/#extras chainlit = ["chainlit"] diff --git a/src/raglite/_database.py b/src/raglite/_database.py index ecef10e..22abad5 100644 --- a/src/raglite/_database.py +++ b/src/raglite/_database.py @@ -6,12 +6,13 @@ from functools import lru_cache from hashlib import sha256 from pathlib import Path -from typing import Any +from typing import Any, cast from xml.sax.saxutils import escape import numpy as np from markdown_it import MarkdownIt from packaging import version +from packaging.version import Version from pydantic import ConfigDict from sqlalchemy.engine import Engine, make_url from sqlmodel import JSON, Column, Field, Relationship, Session, SQLModel, create_engine, text @@ -314,8 +315,10 @@ def from_chunks( @lru_cache(maxsize=1) def _pgvector_version(session: Session) -> Version: try: - result = session.execute(text("SELECT extversion FROM pg_extension WHERE extname = 'vector'")) - pgvector_version = version.parse(result.scalar()) + result = session.execute( + text("SELECT extversion FROM pg_extension WHERE extname = 'vector'") + ) + pgvector_version = version.parse(cast(str, result.scalar_one())) except Exception as e: error_message = "Unable to parse pgvector version, is pgvector installed?" raise ValueError(error_message) from e @@ -372,7 +375,7 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: """ ) ) - base_sql = f""" + create_vector_index_sql = f""" CREATE INDEX IF NOT EXISTS vector_search_chunk_index ON chunk_embedding USING hnsw ( (embedding::halfvec({embedding_dim})) @@ -381,14 +384,10 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: SET hnsw.ef_search = {20 * 4 * 8}; """ # Add iterative scan if version >= 0.8.0 - pgvector_version = _get_pgvector_version(session) - if pgvector_version and version.parse(pgvector_version) >= version.parse("0.8.0"): - sql = f"""{base_sql}; - SET hnsw.iterative_scan = {'relaxed_order' if config.reranker else 'strict_order'}; - """ - else: - sql = f"{base_sql};" - session.execute(text(sql)) + pgvector_version = _pgvector_version(session) + if pgvector_version and pgvector_version >= version.parse("0.8.0"): + create_vector_index_sql += f"\nSET hnsw.iterative_scan = {'relaxed_order' if config.reranker else 'strict_order'};" + session.execute(text(create_vector_index_sql)) session.commit() elif db_backend == "sqlite": # Create a virtual table for keyword search on the chunk table. @@ -396,39 +395,31 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: # [1] https://www.sqlite.org/fts5.html#external_content_tables with Session(engine) as session: session.execute( - text( - """ + text(""" CREATE VIRTUAL TABLE IF NOT EXISTS keyword_search_chunk_index USING fts5(body, content='chunk', content_rowid='rowid'); - """ - ) + """) ) session.execute( - text( - """ + text(""" CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_insert AFTER INSERT ON chunk BEGIN INSERT INTO keyword_search_chunk_index(rowid, body) VALUES (new.rowid, new.body); END; - """ - ) + """) ) session.execute( - text( - """ + text(""" CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_delete AFTER DELETE ON chunk BEGIN INSERT INTO keyword_search_chunk_index(keyword_search_chunk_index, rowid, body) VALUES('delete', old.rowid, old.body); END; - """ - ) + """) ) session.execute( - text( - """ + text(""" CREATE TRIGGER IF NOT EXISTS keyword_search_chunk_index_auto_update AFTER UPDATE ON chunk BEGIN INSERT INTO keyword_search_chunk_index(keyword_search_chunk_index, rowid, body) VALUES('delete', old.rowid, old.body); INSERT INTO keyword_search_chunk_index(rowid, body) VALUES (new.rowid, new.body); END; - """ - ) + """) ) session.commit() return engine From d654872656a0b1a8fa7669380202bad7a0b6a00d Mon Sep 17 00:00:00 2001 From: Manolo Santos Date: Thu, 12 Dec 2024 12:56:36 +0100 Subject: [PATCH 4/5] fix: remove useless caching --- src/raglite/_database.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/raglite/_database.py b/src/raglite/_database.py index 22abad5..c339c82 100644 --- a/src/raglite/_database.py +++ b/src/raglite/_database.py @@ -312,7 +312,6 @@ def from_chunks( ) -@lru_cache(maxsize=1) def _pgvector_version(session: Session) -> Version: try: result = session.execute( From 7b05cb03ce66801fcd3ef5b5d4a6270e64a45c60 Mon Sep 17 00:00:00 2001 From: Laurent Sorber Date: Sun, 15 Dec 2024 11:48:54 +0100 Subject: [PATCH 5/5] chore: relock dependencies --- poetry.lock | 4 ++-- src/raglite/_database.py | 10 ++++------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/poetry.lock b/poetry.lock index b207e92..4429f28 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. [[package]] name = "aiofiles" @@ -6813,4 +6813,4 @@ ragas = ["ragas"] [metadata] lock-version = "2.0" python-versions = ">=3.10,<4.0" -content-hash = "b3a14066711fe4caec356d0aa18514495d44ac253371d2560fc0c5aea890aaef" +content-hash = "239db3c85866a30b063fa6dfe538bbbe92ba659419e47446d5529fbd5bb3831a" diff --git a/src/raglite/_database.py b/src/raglite/_database.py index c339c82..8d6c179 100644 --- a/src/raglite/_database.py +++ b/src/raglite/_database.py @@ -368,11 +368,9 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: with Session(engine) as session: metrics = {"cosine": "cosine", "dot": "ip", "euclidean": "l2", "l1": "l1", "l2": "l2"} session.execute( - text( - """ - CREATE INDEX IF NOT EXISTS keyword_search_chunk_index ON chunk USING GIN (to_tsvector('simple', body)); - """ - ) + text(""" + CREATE INDEX IF NOT EXISTS keyword_search_chunk_index ON chunk USING GIN (to_tsvector('simple', body)); + """) ) create_vector_index_sql = f""" CREATE INDEX IF NOT EXISTS vector_search_chunk_index ON chunk_embedding @@ -382,7 +380,7 @@ def create_database_engine(config: RAGLiteConfig | None = None) -> Engine: ); SET hnsw.ef_search = {20 * 4 * 8}; """ - # Add iterative scan if version >= 0.8.0 + # Enable iterative scan for pgvector v0.8.0 and up. pgvector_version = _pgvector_version(session) if pgvector_version and pgvector_version >= version.parse("0.8.0"): create_vector_index_sql += f"\nSET hnsw.iterative_scan = {'relaxed_order' if config.reranker else 'strict_order'};"