infra: migrate to uv

.github/actions/uv_setup/action.yml

@@ -0,0 +1,23 @@
+# TODO: https://docs.astral.sh/uv/guides/integration/github/#caching
+
+name: uv-install
+description: Set up Python and uv
+
+inputs:
+  python-version:
+    description: Python version, supporting MAJOR.MINOR only
+    required: true
+
+env:
+  UV_VERSION: "0.5.25"
+
+runs:
+  using: composite
+  steps:
+    - uses: actions/checkout@v4
+
+    - name: Install uv and set the python version
+      uses: astral-sh/setup-uv@v5
+      with:
+        version: ${{ env.UV_VERSION }}
+        python-version: ${{ inputs.python-version }}

.github/scripts/check_diff.py

@@ -7,6 +7,8 @@
 from pathlib import Path
 import tomllib
 
+from packaging.requirements import Requirement
+
 from get_min_versions import get_min_version_from_toml
 
 
@@ -61,15 +63,17 @@ def dependents_graph() -> dict:
 
         # load regular and test deps from pyproject.toml
         with open(path, "rb") as f:
-            pyproject = tomllib.load(f)["tool"]["poetry"]
+            pyproject = tomllib.load(f)
 
         pkg_dir = "libs" + "/".join(path.split("libs")[1].split("/")[:-1])
         for dep in [
-            *pyproject["dependencies"].keys(),
-            *pyproject["group"]["test"]["dependencies"].keys(),
+            *pyproject["project"]["dependencies"],
+            *pyproject["dependency-groups"]["test"],
         ]:
+            requirement = Requirement(dep)
+            package_name = requirement.name
             if "langchain" in dep:
-                dependents[dep].add(pkg_dir)
+                dependents[package_name].add(pkg_dir)
                 continue
 
         # load extended deps from extended_testing_deps.txt
@@ -120,8 +124,7 @@ def _get_configs_for_single_dir(job: str, dir_: str) -> List[Dict[str, str]]:
         py_versions = ["3.9", "3.10", "3.11", "3.12", "3.13"]
     # custom logic for specific directories
     elif dir_ == "libs/partners/milvus":
-        # milvus poetry doesn't allow 3.12 because they
-        # declare deps in funny way
+        # milvus doesn't allow 3.12 because they declare deps in funny way
         py_versions = ["3.9", "3.11"]
 
     elif dir_ in PY_312_MAX_PACKAGES:
@@ -148,17 +151,17 @@ def _get_configs_for_single_dir(job: str, dir_: str) -> List[Dict[str, str]]:
 def _get_pydantic_test_configs(
     dir_: str, *, python_version: str = "3.11"
 ) -> List[Dict[str, str]]:
-    with open("./libs/core/poetry.lock", "rb") as f:
-        core_poetry_lock_data = tomllib.load(f)
-    for package in core_poetry_lock_data["package"]:
+    with open("./libs/core/uv.lock", "rb") as f:
+        core_uv_lock_data = tomllib.load(f)
+    for package in core_uv_lock_data["package"]:
         if package["name"] == "pydantic":
             core_max_pydantic_minor = package["version"].split(".")[1]
             break
 
-    with open(f"./{dir_}/poetry.lock", "rb") as f:
-        dir_poetry_lock_data = tomllib.load(f)
+    with open(f"./{dir_}/uv.lock", "rb") as f:
+        dir_uv_lock_data = tomllib.load(f)
 
-    for package in dir_poetry_lock_data["package"]:
+    for package in dir_uv_lock_data["package"]:
         if package["name"] == "pydantic":
             dir_max_pydantic_minor = package["version"].split(".")[1]
             break
@@ -304,7 +307,7 @@ def _get_configs_for_multi_dirs(
                 f"Unknown lib: {file}. check_diff.py likely needs "
                 "an update for this new library!"
             )
-        elif file.startswith("docs/") or file in ["pyproject.toml", "poetry.lock"]: # docs or root poetry files
+        elif file.startswith("docs/") or file in ["pyproject.toml", "uv.lock"]: # docs or root uv files
             docs_edited = True
             dirs_to_run["lint"].add(".")
 

.github/scripts/check_prerelease_dependencies.py

@@ -10,26 +10,25 @@
         toml_data = tomllib.load(file)
 
     # see if we're releasing an rc
-    version = toml_data["tool"]["poetry"]["version"]
+    version = toml_data["project"]["version"]
     releasing_rc = "rc" in version or "dev" in version
 
     # if not, iterate through dependencies and make sure none allow prereleases
     if not releasing_rc:
-        dependencies = toml_data["tool"]["poetry"]["dependencies"]
-        for lib in dependencies:
-            dep_version = dependencies[lib]
+        dependencies = toml_data["project"]["dependencies"]
+        for dep_version in dependencies:
             dep_version_string = (
                 dep_version["version"] if isinstance(dep_version, dict) else dep_version
             )
 
             if "rc" in dep_version_string:
                 raise ValueError(
-                    f"Dependency {lib} has a prerelease version. Please remove this."
+                    f"Dependency {dep_version} has a prerelease version. Please remove this."
                 )
 
             if isinstance(dep_version, dict) and dep_version.get(
                 "allow-prereleases", False
             ):
                 raise ValueError(
-                    f"Dependency {lib} has allow-prereleases set to true. Please remove this."
+                    f"Dependency {dep_version} has allow-prereleases set to true. Please remove this."
                 )

.github/scripts/get_min_versions.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 import sys
 from typing import Optional
 
@@ -7,6 +8,7 @@
     # for python 3.10 and below, which doesnt have stdlib tomllib
     import tomli as tomllib
 
+from packaging.requirements import Requirement
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
@@ -94,6 +96,23 @@ def get_minimum_version(package_name: str, spec_string: str) -> Optional[str]:
     return str(min(valid_versions)) if valid_versions else None
 
 
+def _check_python_version_from_requirement(
+    requirement: Requirement, python_version: str
+) -> bool:
+    if not requirement.marker:
+        return True
+    else:
+        marker_str = str(requirement.marker)
+        if "python_version" or "python_full_version" in marker_str:
+            python_version_str = "".join(
+                char
+                for char in marker_str
+                if char.isdigit() or char in (".", "<", ">", "=", ",")
+            )
+            return check_python_version(python_version, python_version_str)
+        return True
+
+
 def get_min_version_from_toml(
     toml_path: str,
     versions_for: str,
@@ -105,8 +124,10 @@ def get_min_version_from_toml(
     with open(toml_path, "rb") as file:
         toml_data = tomllib.load(file)
 
-    # Get the dependencies from tool.poetry.dependencies
-    dependencies = toml_data["tool"]["poetry"]["dependencies"]
+    dependencies = defaultdict(list)
+    for dep in toml_data["project"]["dependencies"]:
+        requirement = Requirement(dep)
+        dependencies[requirement.name].append(requirement)
 
     # Initialize a dictionary to store the minimum versions
     min_versions = {}
@@ -121,17 +142,11 @@ def get_min_version_from_toml(
         if lib in dependencies:
             if include and lib not in include:
                 continue
-            # Get the version string
-            version_string = dependencies[lib]
-
-            if isinstance(version_string, dict):
-                version_string = version_string["version"]
-            if isinstance(version_string, list):
-                version_string = [
-                    vs
-                    for vs in version_string
-                    if check_python_version(python_version, vs["python"])
-                ][0]["version"]
+            requirements = dependencies[lib]
+            for requirement in requirements:
+                if _check_python_version_from_requirement(requirement, python_version):
+                    version_string = str(requirement.specifier)
+                    break
 
             # Use parse_version to get the minimum supported version from version_string
             min_version = get_minimum_version(lib, version_string)

.github/workflows/_compile_integration_test.yml

@@ -13,7 +13,7 @@ on:
         description: "Python version to use"
 
 env:
-  POETRY_VERSION: "1.8.4"
+  UV_FROZEN: "true"
 
 jobs:
   build:
@@ -22,25 +22,22 @@ jobs:
         working-directory: ${{ inputs.working-directory }}
     runs-on: ubuntu-latest
     timeout-minutes: 20
-    name: "poetry run pytest -m compile tests/integration_tests #${{ inputs.python-version }}"
+    name: "uv run pytest -m compile tests/integration_tests #${{ inputs.python-version }}"
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up Python ${{ inputs.python-version }} + Poetry ${{ env.POETRY_VERSION }}
-        uses: "./.github/actions/poetry_setup"
+      - name: Set up Python ${{ inputs.python-version }} + uv
+        uses: "./.github/actions/uv_setup"
         with:
           python-version: ${{ inputs.python-version }}
-          poetry-version: ${{ env.POETRY_VERSION }}
-          working-directory: ${{ inputs.working-directory }}
-          cache-key: compile-integration
 
       - name: Install integration dependencies
         shell: bash
-        run: poetry install --with=test_integration,test
+        run: uv sync --group test --group test_integration
 
       - name: Check integration tests compile
         shell: bash
-        run: poetry run pytest -m compile tests/integration_tests
+        run: uv run pytest -m compile tests/integration_tests
 
       - name: Ensure the tests did not create any additional files
         shell: bash

.github/workflows/_integration_test.yml

@@ -12,7 +12,7 @@ on:
         description: "Python version to use"
 
 env:
-  POETRY_VERSION: "1.8.4"
+  UV_FROZEN: "true"
 
 jobs:
   build:
@@ -24,22 +24,19 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up Python ${{ inputs.python-version }} + Poetry ${{ env.POETRY_VERSION }}
-        uses: "./.github/actions/poetry_setup"
+      - name: Set up Python ${{ inputs.python-version }} + uv
+        uses: "./.github/actions/uv_setup"
         with:
           python-version: ${{ inputs.python-version }}
-          poetry-version: ${{ env.POETRY_VERSION }}
-          working-directory: ${{ inputs.working-directory }}
-          cache-key: core
 
       - name: Install dependencies
         shell: bash
-        run: poetry install --with test,test_integration
+        run: uv sync --group test --group test_integration
 
       - name: Install deps outside pyproject
         if: ${{ startsWith(inputs.working-directory, 'libs/community/') }}
         shell: bash
-        run: poetry run pip install "boto3<2" "google-cloud-aiplatform<2"
+        run: VIRTUAL_ENV=.venv uv pip install "boto3<2" "google-cloud-aiplatform<2"
 
       - name: Run integration tests
         shell: bash

.github/workflows/_lint.yml

@@ -13,12 +13,13 @@ on:
         description: "Python version to use"
 
 env:
-  POETRY_VERSION: "1.8.4"
   WORKDIR: ${{ inputs.working-directory == '' && '.' || inputs.working-directory }}
 
   # This env var allows us to get inline annotations when ruff has complaints.
   RUFF_OUTPUT_FORMAT: github
 
+  UV_FROZEN: "true"
+
 jobs:
   build:
     name: "make lint #${{ inputs.python-version }}"
@@ -27,25 +28,10 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - name: Set up Python ${{ inputs.python-version }} + Poetry ${{ env.POETRY_VERSION }}
-        uses: "./.github/actions/poetry_setup"
+      - name: Set up Python ${{ inputs.python-version }} + uv
+        uses: "./.github/actions/uv_setup"
         with:
           python-version: ${{ inputs.python-version }}
-          poetry-version: ${{ env.POETRY_VERSION }}
-          working-directory: ${{ inputs.working-directory }}
-          cache-key: lint-with-extras
-
-      - name: Check Poetry File
-        shell: bash
-        working-directory: ${{ inputs.working-directory }}
-        run: |
-          poetry check
-
-      - name: Check lock file
-        shell: bash
-        working-directory: ${{ inputs.working-directory }}
-        run: |
-          poetry lock --check
 
       - name: Install dependencies
         # Also installs dev/lint/test/typing dependencies, to ensure we have
@@ -58,17 +44,7 @@ jobs:
         # It doesn't matter how you change it, any change will cause a cache-bust.
         working-directory: ${{ inputs.working-directory }}
         run: |
-          poetry install --with lint,typing
-
-      - name: Get .mypy_cache to speed up mypy
-        uses: actions/cache@v4
-        env:
-          SEGMENT_DOWNLOAD_TIMEOUT_MIN: "2"
-        with:
-          path: |
-            ${{ env.WORKDIR }}/.mypy_cache
-          key: mypy-lint-${{ runner.os }}-${{ runner.arch }}-py${{ inputs.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/poetry.lock', inputs.working-directory)) }}
-
+          uv sync --group lint --group typing
 
       - name: Analysing the code with our lint
         working-directory: ${{ inputs.working-directory }}
@@ -87,21 +63,12 @@ jobs:
         if: ${{ ! startsWith(inputs.working-directory, 'libs/partners/') }}
         working-directory: ${{ inputs.working-directory }}
         run: |
-          poetry install --with test
+          uv sync --group test
       - name: Install unit+integration test dependencies
         if: ${{ startsWith(inputs.working-directory, 'libs/partners/') }}
         working-directory: ${{ inputs.working-directory }}
         run: |
-          poetry install --with test,test_integration
-
-      - name: Get .mypy_cache_test to speed up mypy
-        uses: actions/cache@v4
-        env:
-          SEGMENT_DOWNLOAD_TIMEOUT_MIN: "2"
-        with:
-          path: |
-            ${{ env.WORKDIR }}/.mypy_cache_test
-          key: mypy-test-${{ runner.os }}-${{ runner.arch }}-py${{ inputs.python-version }}-${{ inputs.working-directory }}-${{ hashFiles(format('{0}/poetry.lock', inputs.working-directory)) }}
+          uv sync --group test --group test_integration
 
       - name: Analysing the code with our lint
         working-directory: ${{ inputs.working-directory }}