feat: add PostgreSQL driver

.cruft.json

@@ -1,15 +1,15 @@
 {
-  "template": "https://github.com/radix-ai/poetry-cookiecutter",
+  "template": "https://github.com/superlinear-ai/poetry-cookiecutter",
   "commit": "a969f1d182ec39d7d27ccb1116cf60ba736adcfa",
   "checkout": null,
   "context": {
     "cookiecutter": {
       "project_type": "package",
       "project_name": "RAGLite",
-      "project_description": "A RAG extension for SQLite.",
-      "project_url": "https://github.com/radix-ai/raglite",
+      "project_description": "A Python package for Retrieval-Augmented Generation (RAG) with SQLite or PostgreSQL.",
+      "project_url": "https://github.com/superlinear-ai/raglite",
       "author_name": "Laurent Sorber",
-      "author_email": "laurent@radix.ai",
+      "author_email": "laurent@superlinear.eu",
       "python_version": "3.10",
       "development_environment": "strict",
       "with_conventional_commits": "1",
@@ -22,8 +22,8 @@
       "__docstring_style": "NumPy",
       "__project_name_kebab_case": "raglite",
       "__project_name_snake_case": "raglite",
-      "_template": "https://github.com/radix-ai/poetry-cookiecutter"
+      "_template": "https://github.com/superlinear-ai/poetry-cookiecutter"
     }
   },
   "directory": null
-}
+}

README.md

@@ -1,22 +1,23 @@
-[![Open in Dev Containers](https://img.shields.io/static/v1?label=Dev%20Containers&message=Open&color=blue&logo=visualstudiocode)](https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/radix-ai/raglite) [![Open in GitHub Codespaces](https://img.shields.io/static/v1?label=GitHub%20Codespaces&message=Open&color=blue&logo=github)](https://github.com/codespaces/new?hide_repo_select=true&ref=main&repo=812973394&skip_quickstart=true)
+[![Open in Dev Containers](https://img.shields.io/static/v1?label=Dev%20Containers&message=Open&color=blue&logo=visualstudiocode)](https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/superlinear-ai/raglite) [![Open in GitHub Codespaces](https://img.shields.io/static/v1?label=GitHub%20Codespaces&message=Open&color=blue&logo=github)](https://github.com/codespaces/new?hide_repo_select=true&ref=main&repo=812973394&skip_quickstart=true)
 
 # 🧵 RAGLite
 
-RAGLite is a Python package for Retrieval-Augmented Generation (RAG) with SQLite.
+RAGLite is a Python package for Retrieval-Augmented Generation (RAG) with PostgreSQL or SQLite.
 
 ## Features
 
 1. ❤️ Only lightweight and permissive open source dependencies (e.g., no [PyTorch](https://github.com/pytorch/pytorch), [LangChain](https://github.com/langchain-ai/langchain), or [PyMuPDF](https://github.com/pymupdf/PyMuPDF))
-2. 🔒 Fully local RAG with [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) as an LLM provider and [SQLite](https://github.com/sqlite/sqlite) as a local database
-3. 🚀 Acceleration with Metal on macOS and with CUDA on Linux and Windows
-4. 📖 PDF to Markdown conversion on top of [pdftext](https://github.com/VikParuchuri/pdftext) and [pypdfium2](https://github.com/pypdfium2-team/pypdfium2)
-5. ✂️ Optimal [level 4 semantic chunking](https://medium.com/@anuragmishra_27746/five-levels-of-chunking-strategies-in-rag-notes-from-gregs-video-7b735895694d) by solving a [binary integer programming problem](https://en.wikipedia.org/wiki/Integer_programming)
-6. 📌 Markdown-based [contextual chunk headings](https://d-star.ai/solving-the-out-of-context-chunk-problem-for-rag)
-7. 🌈 Combined sentence-level and chunk-level matching with [multi-vector chunk retrieval](https://python.langchain.com/v0.2/docs/how_to/multi_vector/)
-8. 🌀 Optimal [closed-form linear query adapter](src/raglite/_query_adapter.py) by solving an [orthogonal Procrustes problem](https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem)
-9. 🔍 [Hybrid search](https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf) that combines [SQLite's BM25 full-text search](https://sqlite.org/fts5.html) with [PyNNDescent's ANN vector search](https://github.com/lmcinnes/pynndescent)
-10. ✍️ Optional support for conversion of any input document to Markdown with [Pandoc](https://github.com/jgm/pandoc)
-11. ✅ Optional support for evaluation of retrieval and generation with [Ragas](https://github.com/explodinggradients/ragas)
+2. 🧠 Your choice of local LLM with [llama-cpp-python](https://github.com/abetlen/llama-cpp-python)
+3. 💾 Your choice of [PostgreSQL](https://github.com/postgres/postgres) or [SQLite](https://github.com/sqlite/sqlite) as a full-text & vector search database
+4. 🚀 Acceleration with Metal on macOS and with CUDA on Linux and Windows
+5. 📖 PDF to Markdown conversion on top of [pdftext](https://github.com/VikParuchuri/pdftext) and [pypdfium2](https://github.com/pypdfium2-team/pypdfium2)
+6. ✂️ Optimal [level 4 semantic chunking](https://medium.com/@anuragmishra_27746/five-levels-of-chunking-strategies-in-rag-notes-from-gregs-video-7b735895694d) by solving a [binary integer programming problem](https://en.wikipedia.org/wiki/Integer_programming)
+7. 📌 Markdown-based [contextual chunk headings](https://d-star.ai/solving-the-out-of-context-chunk-problem-for-rag)
+8. 🌈 Combined sentence-level and chunk-level matching with [multi-vector chunk retrieval](https://python.langchain.com/v0.2/docs/how_to/multi_vector/)
+9. 🌀 Optimal [closed-form linear query adapter](src/raglite/_query_adapter.py) by solving an [orthogonal Procrustes problem](https://en.wikipedia.org/wiki/Orthogonal_Procrustes_problem)
+10. 🔍 [Hybrid search](https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf) that combines the database's built-in full-text search ([tsvector](https://www.postgresql.org/docs/current/datatype-textsearch.html) in PostgreSQL, [FTS5](https://www.sqlite.org/fts5.html) in SQLite) with their native vector search extensions ([pgvector](https://github.com/pgvector/pgvector) in PostgreSQL, [sqlite-vec](https://github.com/asg017/sqlite-vec) in SQLite)
+11. ✍️ Optional support for conversion of any input document to Markdown with [Pandoc](https://github.com/jgm/pandoc)
+12. ✅ Optional support for evaluation of retrieval and generation performance with [Ragas](https://github.com/explodinggradients/ragas)
 
 ## Installing
 
@@ -145,7 +146,7 @@ evaluation_df = evaluate(answered_evals_df, config=my_config)
 The following development environments are supported:
 
 1. ⭐️ _GitHub Codespaces_: click on _Code_ and select _Create codespace_ to start a Dev Container with [GitHub Codespaces](https://github.com/features/codespaces).
-1. ⭐️ _Dev Container (with container volume)_: click on [Open in Dev Containers](https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/radix-ai/raglite) to clone this repository in a container volume and create a Dev Container with VS Code.
+1. ⭐️ _Dev Container (with container volume)_: click on [Open in Dev Containers](https://vscode.dev/redirect?url=vscode://ms-vscode-remote.remote-containers/cloneInVolume?url=https://github.com/superlinear-ai/raglite) to clone this repository in a container volume and create a Dev Container with VS Code.
 1. _Dev Container_: clone this repository, open it with VS Code, and run <kbd>Ctrl/⌘</kbd> + <kbd>⇧</kbd> + <kbd>P</kbd> → _Dev Containers: Reopen in Container_.
 1. _PyCharm_: clone this repository, open it with PyCharm, and [configure Docker Compose as a remote interpreter](https://www.jetbrains.com/help/pycharm/using-docker-compose-as-a-remote-interpreter.html#docker-compose-remote) with the `dev` service.
 1. _Terminal_: clone this repository, open it with your terminal, and run `docker compose up --detach dev` to start a Dev Container in the background, and then run `docker compose exec dev zsh` to open a shell prompt in the Dev Container.

docker-compose.yml

@@ -12,6 +12,10 @@ services:
         GID: ${GID:-1000}
     environment:
       - POETRY_PYPI_TOKEN_PYPI
+    depends_on:
+      - postgres
+    networks:
+      - raglite-network
     volumes:
       - ..:/workspaces
       - command-history-volume:/home/user/.history/
@@ -21,21 +25,34 @@ services:
     stdin_open: true
     tty: true
     entrypoint: []
-    command:
-      [
-        "sh",
-        "-c",
-        "sudo chown user $$SSH_AUTH_SOCK && cp --update /opt/build/poetry/poetry.lock /workspaces/raglite/ && mkdir -p /workspaces/raglite/.git/hooks/ && cp --update /opt/build/git/* /workspaces/raglite/.git/hooks/ && zsh"
-      ]
+    command: [ "sh", "-c", "sudo chown user $$SSH_AUTH_SOCK && cp --update /opt/build/poetry/poetry.lock /workspaces/raglite/ && mkdir -p /workspaces/raglite/.git/hooks/ && cp --update /opt/build/git/* /workspaces/raglite/.git/hooks/ && zsh" ]
     environment:
       - POETRY_PYPI_TOKEN_PYPI
       - SSH_AUTH_SOCK=/run/host-services/ssh-auth.sock
+    depends_on:
+      - postgres
+    networks:
+      - raglite-network
     volumes:
       - ~/.gitconfig:/etc/gitconfig
       - ~/.ssh/known_hosts:/home/user/.ssh/known_hosts
       - ${SSH_AGENT_AUTH_SOCK:-/run/host-services/ssh-auth.sock}:/run/host-services/ssh-auth.sock
     profiles:
       - dev
 
+  postgres:
+    image: pgvector/pgvector:pg16
+    environment:
+      POSTGRES_USER: raglite_user
+      POSTGRES_PASSWORD: raglite_password
+    networks:
+      - raglite-network
+    tmpfs:
+      - /var/lib/postgresql/data
+
+networks:
+  raglite-network:
+    driver: bridge
+
 volumes:
   command-history-volume:

pyproject.toml

@@ -5,10 +5,10 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry] # https://python-poetry.org/docs/pyproject/
 name = "raglite"
 version = "0.0.0"
-description = "A Python package for Retrieval-Augmented Generation (RAG) with SQLite."
-authors = ["Laurent Sorber <laurent@radix.ai>"]
+description = "A Python package for Retrieval-Augmented Generation (RAG) with SQLite or PostgreSQL."
+authors = ["Laurent Sorber <laurent@superlinear.eu>"]
 readme = "README.md"
-repository = "https://github.com/radix-ai/raglite"
+repository = "https://github.com/superlinear-ai/raglite"
 
 [tool.commitizen] # https://commitizen-tools.github.io/commitizen/config/
 bump_message = "bump(release): v$current_version → v$new_version"
@@ -51,6 +51,7 @@ pydantic = ">=2.7.0"
 # Approximate Nearest Neighbors:
 pynndescent = ">=0.5.12"
 # Storage:
+pg8000 = ">=1.31.2"
 sqlmodel-slim = ">=0.0.18"
 # Progress:
 tqdm = ">=4.66.0"
@@ -114,7 +115,7 @@ warn_unreachable = true
 
 [tool.pytest.ini_options] # https://docs.pytest.org/en/latest/reference/reference.html#ini-options-ref
 addopts = "--color=yes --doctest-modules --exitfirst --failed-first --strict-config --strict-markers --verbosity=2 --junitxml=reports/pytest.xml"
-filterwarnings = ["error", "ignore::DeprecationWarning"]
+filterwarnings = ["error", "ignore::DeprecationWarning", "ignore::pytest.PytestUnraisableExceptionWarning"]
 testpaths = ["src", "tests"]
 xfail_strict = true
 

src/raglite/__init__.py

@@ -2,7 +2,7 @@
 
 from raglite._config import RAGLiteConfig
 from raglite._eval import answer_evals, evaluate, insert_evals
-from raglite._index import insert_document, update_vector_index
+from raglite._index import insert_document
 from raglite._query_adapter import update_query_adapter
 from raglite._rag import rag
 from raglite._search import (
@@ -18,7 +18,6 @@
     "RAGLiteConfig",
     # Index
     "insert_document",
-    "update_vector_index",
     # Search
     "fusion_search",
     "hybrid_search",

src/raglite/_config.py

@@ -3,23 +3,21 @@
 from dataclasses import dataclass, field
 from functools import lru_cache
 
-import numpy as np
-import numpy.typing as npt
 from llama_cpp import Llama, LlamaRAMCache, llama_supports_gpu_offload  # type: ignore[attr-defined]
 from sqlalchemy.engine import URL
 
 
 @lru_cache(maxsize=1)
 def default_llm() -> Llama:
     """Get default LLM."""
-    # Select the best available LLM for the given accelerator.
+    # Select the best available LLM for the given accelerator:
+    # - Llama-3.1-8B-instruct on GPU.
+    # - Phi-3.5-mini-instruct on CPU.
     if llama_supports_gpu_offload():
-        # Llama-3.1-8B-instruct on GPU.
         repo_id = "bartowski/Meta-Llama-3.1-8B-Instruct-GGUF"  # https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct
         filename = "*Q4_K_M.gguf"
         n_ctx = 8192
     else:
-        # Phi-3.1-mini-128k-instruct on CPU.
         repo_id = "bartowski/Phi-3.5-mini-instruct-GGUF"  # https://huggingface.co/microsoft/Phi-3.5-mini-instruct
         filename = "*Q4_K_M.gguf"
         n_ctx = 4096
@@ -61,7 +59,6 @@ class RAGLiteConfig:
     # Embedder config used for indexing.
     embedder: Llama = field(default_factory=default_embedder)
     embedder_batch_size: int = 128
-    embedder_dtype: npt.DTypeLike = np.float16
     embedder_normalize: bool = True
     sentence_embedding_weight: float = 0.5  # Between 0 (chunk level) and 1 (sentence level).
     # Chunker config used to partition documents into chunks.
@@ -70,7 +67,5 @@ class RAGLiteConfig:
     # Database config.
     db_url: str | URL = "sqlite:///raglite.sqlite"
     # Vector search config.
-    vector_search_index_id: str = "default"
     vector_search_index_metric: str = "cosine"  # The query adapter supports "dot" and "cosine".
-    # Query adapter config.
-    enable_query_adapter: bool = True
+    vector_search_query_adapter: bool = True