depot/third_party/nixpkgs/pkgs/development/python-modules/manifest-ml/default.nix

{
  lib,
  accelerate,
  aiohttp,
  buildPythonPackage,
  fastapi,
  fetchFromGitHub,
  flask,
  numpy,
  pg8000,
  pillow,
  pydantic,
  pytestCheckHook,
  pythonOlder,
  pythonRelaxDepsHook,
  redis,
  requests,
  sentence-transformers,
  setuptools,
  sqlalchemy,
  sqlitedict,
  tenacity,
  tiktoken,
  torch,
  transformers,
  uvicorn,
  xxhash,
}:

buildPythonPackage rec {
  pname = "manifest-ml";
  version = "0.1.9";
  pyproject = true;

  disabled = pythonOlder "3.8";

  src = fetchFromGitHub {
    owner = "HazyResearch";
    repo = "manifest";
    rev = "refs/tags/v${version}";
    hash = "sha256-6m1XZOXzflBYyq9+PinbrW+zqvNGFN/aRDHH1b2Me5E=";
  };

  __darwinAllowLocalNetworking = true;

  pythonRelaxDeps = [ "pydantic" ];

  build-system = [
    pythonRelaxDepsHook
    setuptools
  ];

  dependencies = [
    numpy
    pydantic
    redis
    requests
    aiohttp
    sqlitedict
    tenacity
    tiktoken
    xxhash
  ];

  passthru.optional-dependencies = {
    api = [
      accelerate
      # deepspeed
      # diffusers
      flask
      sentence-transformers
      torch
      transformers
    ];
    app = [
      fastapi
      uvicorn
    ];
    diffusers = [ pillow ];
    gcp = [
      pg8000
      # cloud-sql-python-connector
      sqlalchemy
    ];
  };

  nativeCheckInputs = [
    pytestCheckHook
  ] ++ lib.flatten (lib.attrValues passthru.optional-dependencies);

  preCheck = ''
    export HOME=$TMPDIR
  '';

  pytestFlagsArray = [
    # this file tries importing `deepspeed`, which is not yet packaged in nixpkgs
    "--ignore=tests/test_huggingface_api.py"
  ];

  disabledTests = [
    # Tests require DB access
    "test_init"
    "test_key_get_and_set"
    "test_get"
    # Tests require network access
    "test_abatch_run"
    "test_batch_run"
    "test_retry_handling"
    "test_run_chat"
    "test_run"
    "test_score_run"
    # Test is time-senstive
    "test_timing"
  ];

  pythonImportsCheck = [ "manifest" ];

  meta = with lib; {
    description = "Manifest for Prompting Foundation Models";
    homepage = "https://github.com/HazyResearch/manifest";
    changelog = "https://github.com/HazyResearch/manifest/releases/tag/v${version}";
    license = licenses.asl20;
    maintainers = with maintainers; [ natsukium ];
  };
}