From 2610733f3f12403e0f4805f75037235c281e41f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 26 Nov 2024 13:20:56 +0000
Subject: [PATCH] Add a primitive script to generate Poetry commands to sync
 with Nix

This is not fully automated, since getting the Nix versions may be
unresolvable. However, it does take most of the work out of doing
this manually.
---
 flake.nix                 |  2 ++
 server/bounds-from-nix.py | 40 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)
 create mode 100755 server/bounds-from-nix.py

diff --git a/flake.nix b/flake.nix
index f26a983e..83cedfa6 100644
--- a/flake.nix
+++ b/flake.nix
@@ -148,6 +148,8 @@
         };
 
         packages = rec {
+          inherit server;
+
           default = pkgs.writeShellApplication {
             name = "text-generation-inference";
             runtimeInputs = [
diff --git a/server/bounds-from-nix.py b/server/bounds-from-nix.py
new file mode 100755
index 00000000..42422b8b
--- /dev/null
+++ b/server/bounds-from-nix.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+import json
+import subprocess
+from typing import Dict, Union
+import toml
+
+# Special cases that have download URLs.
+SKIP = {"attention-kernels", "marlin-kernels", "moe-kernels"}
+
+
+def is_optional(info: Union[str, Dict[str, str]]) -> bool:
+    return isinstance(info, dict) and "optional" in info and info["optional"]
+
+
+if __name__ == "__main__":
+    with open("pyproject.toml") as f:
+        pyproject = toml.load(f)
+
+    nix_packages = json.loads(
+        subprocess.run(
+            ["nix", "develop", ".#server", "--command", "pip", "list", "--format=json"],
+            stdout=subprocess.PIPE,
+        ).stdout
+    )
+
+    nix_packages = {pkg["name"]: pkg["version"] for pkg in nix_packages}
+
+    packages = []
+    optional_packages = []
+
+    for package, info in pyproject["tool"]["poetry"]["dependencies"].items():
+        if package in nix_packages and package not in SKIP:
+            if is_optional(info):
+                optional_packages.append(f'"{package}@^{nix_packages[package]}"')
+            else:
+                packages.append(f'"{package}@^{nix_packages[package]}"')
+
+    print(f"poetry add {' '.join(packages)}")
+    print(f"poetry add --optional {' '.join(optional_packages)}")