From c625352b5326f9b52d20b2a6c6d9bd1c5aa6bfe3 Mon Sep 17 00:00:00 2001
From: happysalada <raphael@megzari.com>
Date: Wed, 12 Jul 2023 09:49:00 +0800
Subject: [PATCH] python310Packages.argilla: 1.8.0 -> 1.12.0

---
 .../python-modules/argilla/default.nix        | 171 ++++++++++++------
 1 file changed, 119 insertions(+), 52 deletions(-)

diff --git a/pkgs/development/python-modules/argilla/default.nix b/pkgs/development/python-modules/argilla/default.nix
index ee1f6c7de6cc..c04310d28b0d 100644
--- a/pkgs/development/python-modules/argilla/default.nix
+++ b/pkgs/development/python-modules/argilla/default.nix
@@ -2,6 +2,7 @@
 , buildPythonPackage
 , fetchFromGitHub
 , pythonOlder
+, pythonRelaxDepsHook
 , deprecated
 , rich
 , backoff
@@ -14,33 +15,106 @@
 , httpx
 , pandas
 , monotonic
-# test dependencies
-, pytestCheckHook
+# optional-dependencies
 , fastapi
-, sqlalchemy
 , opensearch-py
-, factory_boy
 , elasticsearch8
-, elastic-transport
-, luqum
-, pytest-asyncio
-, passlib
-, python-jose
-, alembic
 , uvicorn
+, smart-open
+, brotli-asgi
+, alembic
+, sqlalchemy
+, greenlet
+, aiosqlite
+, luqum
+, scikit-learn
+, aiofiles
+, pyyaml
+, python-multipart
+, python-jose
+, passlib
+, psutil
+# , segment-analytics-python
+, asyncpg
+, psycopg2
 , schedule
 , prodict
-, datasets
-, psutil
-, spacy
 , cleanlab
-, snorkel
-, transformers
+, datasets
+, huggingface-hub
+# , flair
 , faiss
+, flyingsquid
+, pgmpy
+, plotly
+, snorkel
+, spacy
+, transformers
+, evaluate
+, seqeval
+# , setfit
+# , span_marker
+, openai
+, peft
+# test dependencies
+, pytestCheckHook
+, pytest-cov
+, pytest-mock
+, pytest-asyncio
+, factory_boy
 }:
 let
   pname = "argilla";
-  version = "1.8.0";
+  version = "1.12.0";
+  optional-dependencies = {
+    server = [
+      fastapi
+      opensearch-py
+      elasticsearch8
+      uvicorn
+      smart-open
+      brotli-asgi
+      alembic
+      sqlalchemy
+      greenlet
+      aiosqlite
+      luqum
+      scikit-learn
+      aiofiles
+      pyyaml
+      python-multipart
+      python-jose
+      passlib
+      psutil
+      # segment-analytics-python
+    ] ++
+      elasticsearch8.optional-dependencies.async ++
+      uvicorn.optional-dependencies.standard ++
+      python-jose.optional-dependencies.cryptography ++
+      passlib.optional-dependencies.bcrypt;
+    postgresql = [ asyncpg psycopg2 ];
+    listeners = [ schedule prodict ];
+    integrations = [
+      pyyaml
+      cleanlab
+      datasets
+      huggingface-hub
+      # flair
+      faiss
+      flyingsquid
+      pgmpy
+      plotly
+      snorkel
+      spacy
+      transformers
+      evaluate
+      seqeval
+      # setfit
+      # span_marker
+      openai
+      peft
+    ] ++ transformers.optional-dependencies.torch;
+  };
 in
 buildPythonPackage {
   inherit pname version;
@@ -52,64 +126,57 @@ buildPythonPackage {
     owner = "argilla-io";
     repo = pname;
     rev = "v${version}";
-    hash = "sha256-pUfuwA/+fe1VVWyGxEkvSuJLNxw3sHmp8cQZecW8GWY=";
+    hash = "sha256-NImtS2bbCfbhbrw12xhGdZp/JVfrB6cHnUHYX3xJ7tw=";
   };
 
-  postPatch = ''
-    substituteInPlace pyproject.toml \
-      --replace '"rich <= 13.0.1"' '"rich"' \
-      --replace '"numpy < 1.24.0"' '"numpy"'
-  '';
+  pythonRelaxDeps = [
+    "typer"
+    "rich"
+    "numpy"
+  ];
+
+  nativeBuildInputs = [
+    pythonRelaxDepsHook
+  ];
 
   propagatedBuildInputs = [
+    httpx
     deprecated
-    rich
-    backoff
     packaging
+    pandas
     pydantic
-    typer
-    tqdm
     wrapt
     numpy
-    pandas
-    httpx
+    tqdm
+    backoff
     monotonic
+    rich
+    typer
   ];
 
+  # still quite a bit of optional dependencies missing
+  doCheck = false;
+
   preCheck = ''
     export HOME=$(mktemp -d)
   '';
 
-  # tests require an opensearch instance running and flyingsquid to be packaged
-  doCheck = false;
-
   nativeCheckInputs = [
     pytestCheckHook
-    fastapi
-    sqlalchemy
-    opensearch-py
-    factory_boy
-    elasticsearch8
-    elastic-transport
-    luqum
+    pytest-cov
+    pytest-mock
     pytest-asyncio
-    passlib
-    python-jose
-    alembic
-    uvicorn
-    schedule
-    prodict
-    datasets
-    psutil
-    spacy
-    cleanlab
-    snorkel
-    transformers
-    faiss
-  ] ++ opensearch-py.optional-dependencies.async;
+    factory_boy
+  ]
+    ++ optional-dependencies.server
+    ++ optional-dependencies.postgresql
+    ++ optional-dependencies.listeners
+    ++ optional-dependencies.integrations;
 
   pytestFlagsArray = [ "--ignore=tests/server/datasets/test_dao.py" ];
 
+  passthru.optional-dependencies = optional-dependencies;
+
   meta = with lib; {
     description = "Argilla: the open-source data curation platform for LLMs";
     homepage = "https://github.com/argilla-io/argilla";