chore: lock transitive dependencies using uv (#371)

This change switches to using uv to manage our dependencies. This is
primarily motivated by the desire to lock the versions of our transitive
dependencies. Also uv is fast.

uv tracks locked dependencies in uv.lock.

In order to not impose uv on end-users who want to install the extension,
we provide a locked requirements file requirements-lock.txt which is used
in the installation process. This file must be updated when uv.lock is
updated, with uv export -o requirements-lock.txt.

Co-authored-by: Jascha <[email protected]>

.github/workflows/ci.yml

@@ -53,6 +53,9 @@ jobs:
 
       - name: Check Python Formatting
         run: docker exec pgai-ext just format-py
+
+      - name: Compare requirements file
+        run: docker exec pgai-ext just check-requirements
 
       - name: Install extension
         run: docker exec pgai-ext just install

DEVELOPMENT.md

@@ -46,6 +46,7 @@ To make changes to the pgai extension, do the following in your developer enviro
    * [Python3](https://www.python.org/downloads/).
    * [Docker](https://docs.docker.com/get-docker/).
    * A Postgres client like [psql](https://www.timescale.com/blog/how-to-install-psql-on-mac-ubuntu-debian-windows/) or [PopSQL](https://docs.timescale.com/use-timescale/latest/popsql/).
+   * [uv](https://docs.astral.sh/uv/getting-started/installation/)
 * Retrieve the API Keys for the LLM cloud providers you are working with.
 * Clone the pgai source:
    ```bash
@@ -166,6 +167,25 @@ To set up the tests:
 
    Providing these keys automatically enables the corresponding tests.
 
+### Dependency management in the pgai extension
+
+The pgai extension uses `uv` for dependency management. Dependencies are listed
+in the [pyproject.toml](./projects/extension/pyproject.toml) file.
+uv is already installed in the development docker container. To use it in your
+local environment, follow uv's [installation instructions](https://github.com/astral-sh/uv#installation).
+
+1. The basic commands are:
+   1. **Install dependencies**: `uv sync`
+   2. **Add a new dependency**: `uv add <package-name>`
+   3. **Remove a dependency**: `uv remove <package-name>`
+
+1. Locking dependencies:
+   To allow for installation of the pgai extension without having uv available,
+   we provide a `requirements-lock.txt` file in addition to the (canonical)
+   `uv.lock` file. Setuptools uses the `requirements-lock.txt` file to install
+   dependencies without uv. This file is automatically generated from the
+   `uv.lock` file with:
+   `uv export --format requirements-txt -o requirements-lock.txt`.
 
 ### The pgai extension architecture
 

projects/extension/Dockerfile

@@ -47,7 +47,7 @@ RUN set -e; \
 FROM base AS pgai-test-db
 ENV PG_MAJOR=${PG_MAJOR}
 ENV PIP_BREAK_SYSTEM_PACKAGES=1
-RUN pip install uv==0.5.9
+RUN pip install uv==0.5.20
 WORKDIR /pgai
 COPY . .
 RUN just build install
@@ -59,17 +59,18 @@ FROM base
 ENV WHERE_AM_I=docker
 USER root
 
-RUN pip install --break-system-packages uv==0.5.9
+RUN pip install --break-system-packages uv==0.5.20
 
 # install pgspot
 RUN set -eux; \
     git clone https://github.com/timescale/pgspot.git /build/pgspot; \
     uv pip install --system --break-system-packages /build/pgspot; \
     rm -rf /build/pgspot
 
-# install our test python dependencies
-COPY requirements-test.txt /build/requirements-test.txt
-RUN uv pip install --system --break-system-packages -r /build/requirements-test.txt
-RUN rm -r /build
-
 WORKDIR /pgai
+
+## install our test python dependencies
+COPY pyproject.toml uv.lock requirements-lock.txt build.py /pgai
+COPY ai /pgai/ai
+RUN uv sync
+ENV PATH="/pgai/.venv/bin:$PATH"

projects/extension/build.py

@@ -656,7 +656,7 @@ def test_server() -> None:
             """
         subprocess.run(cmd, shell=True, check=True, env=os.environ, cwd=ext_dir())
     else:
-        cmd = "fastapi dev server.py"
+        cmd = "uv run fastapi dev server.py"
         subprocess.run(
             cmd,
             shell=True,
@@ -667,7 +667,9 @@ def test_server() -> None:
 
 
 def test() -> None:
-    subprocess.run("pytest", shell=True, check=True, env=os.environ, cwd=tests_dir())
+    subprocess.run(
+        "uv run pytest", shell=True, check=True, env=os.environ, cwd=tests_dir()
+    )
 
 
 def lint_sql() -> None:
@@ -683,7 +685,9 @@ def lint_sql() -> None:
 
 
 def lint_py() -> None:
-    subprocess.run(f"ruff check {ext_dir()}", shell=True, check=True, env=os.environ)
+    subprocess.run(
+        f"uv run ruff check {ext_dir()}", shell=True, check=True, env=os.environ
+    )
 
 
 def lint() -> None:
@@ -693,14 +697,57 @@ def lint() -> None:
 
 def format_py() -> None:
     subprocess.run(
-        f"ruff format --diff {ext_dir()}", shell=True, check=True, env=os.environ
+        f"uv run ruff format --diff {ext_dir()}", shell=True, check=True, env=os.environ
     )
 
 
 def reformat_py() -> None:
     subprocess.run(f"ruff format {ext_dir()}", shell=True, check=True, env=os.environ)
 
 
+def check_requirements() -> None:
+    """
+    Verifies that requirements-lock.txt is up to date with pyproject.toml.
+    Creates a temporary file with the current state and compares it with the existing lock file.
+    """
+    if shutil.which("uv") is None:
+        fatal("uv not found")
+
+    # Create a temporary file to store current requirements
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".txt") as tmp_file:
+        # Generate current requirements
+        subprocess.run(
+            f"uv export --quiet --format requirements-txt -o {tmp_file.name}",
+            shell=True,
+            check=True,
+            env=os.environ,
+            text=True,
+        )
+
+        # Read both files
+        lock_file = ext_dir() / "requirements-lock.txt"
+        if not lock_file.exists():
+            fatal(
+                "requirements-lock.txt does not exist. Run 'uv export --format requirements-txt -o requirements-lock.txt' to create it."
+            )
+
+        from difflib import unified_diff
+
+        with open(lock_file, "r") as f1, open(tmp_file.name, "r") as f2:
+            # Skip the first 3 lines when reading both files since the contain a line with the file name
+            # which will always be different
+            lock_contents = f1.readlines()[3:]
+            current_contents = f2.readlines()[3:]
+
+            diff = list(unified_diff(lock_contents, current_contents))
+            if diff:
+                fatal(
+                    "requirements-lock.txt is out of sync with uv.lock.\n"
+                    "Run 'uv export --format requirements-txt -o requirements-lock.txt' to update it.\n"
+                    + "".join(diff)
+                )
+
+
 def docker_build() -> None:
     subprocess.run(
         f"""docker build --build-arg PG_MAJOR={pg_major()} -t pgai-ext .""",
@@ -721,6 +768,7 @@ def docker_run() -> None:
             "docker run -d --name pgai-ext --hostname pgai-ext -e POSTGRES_HOST_AUTH_METHOD=trust",
             networking,
             f"--mount type=bind,src={ext_dir()},dst=/pgai",
+            "--mount type=volume,dst=/pgai/.venv",
             "-e TEST_ENV_SECRET=super_secret",
             "pgai-ext",
             "-c shared_preload_libraries='timescaledb, pgextwlist'",
@@ -812,6 +860,8 @@ def run() -> None:
             format_py()
         elif action == "reformat-py":
             reformat_py()
+        elif action == "check-requirements":
+            check_requirements()
         elif action == "docker-build":
             docker_build()
         elif action == "docker-run":

projects/extension/justfile

@@ -86,6 +86,9 @@ format-py:
 reformat-py:
 	@./build.py reformat-py
 
+check-requirements:
+    @./build.py check-requirements
+
 docker-build:
 	@PG_MAJOR={{PG_MAJOR}} ./build.py docker-build
 

projects/extension/pyproject.toml

@@ -4,12 +4,34 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pgai"
-dynamic = ["version", "dependencies"]
+dynamic = ["version"]
 requires-python = ">=3.10"
 
+dependencies= [
+    "openai==1.44.0",
+    "tiktoken==0.7.0",
+    "ollama==0.4.5",
+    "anthropic==0.29.0",
+    "cohere==5.5.8",
+    "backoff==2.2.1",
+    "voyageai==0.3.1",
+    "datasets==3.1.0",
+]
+
+[dependency-groups]
+dev = [
+    "ruff==0.6.9",
+    "pytest==8.3.2",
+    "python-dotenv==1.0.1",
+    "fastapi==0.112.0",
+    "fastapi-cli==0.0.5",
+    "psycopg[binary]==3.2.1",
+    "uv==0.5.20",
+]
+
 [tool.setuptools.dynamic]
 version = {attr = "ai.__version__"}
-dependencies = {file = ["requirements.txt"]}
+dependencies = {file = ["requirements-lock.txt"]}
 
 [tool.setuptools]
 packages = ["ai"]