python3Packages.litellm: fix litellm_enterprise module not found

Author: codgician

pkgs/development/python-modules/litellm/default.nix

@@ -5,6 +5,7 @@
   azure-identity,
   azure-keyvault-secrets,
   backoff,
+  boto3,
   buildPythonPackage,
   click,
   cryptography,
@@ -25,17 +26,20 @@
   pydantic,
   pyjwt,
   pynacl,
+  python,
   python-dotenv,
   python-multipart,
   pythonOlder,
   pyyaml,
   requests,
   resend,
+  rich,
   rq,
   tiktoken,
   tokenizers,
   uvloop,
   uvicorn,
+  websockets,
   nixosTests,
   nix-update-script,
 }:
@@ -63,7 +67,6 @@ buildPythonPackage rec {
     importlib-metadata
     jinja2
     jsonschema
-    mcp
     openai
     pydantic
     python-dotenv
@@ -76,36 +79,51 @@ buildPythonPackage rec {
     proxy = [
       apscheduler
       backoff
+      boto3
       cryptography
       fastapi
       fastapi-sso
       gunicorn
+      mcp
       orjson
       pyjwt
+      pynacl
       python-multipart
       pyyaml
+      rich
       rq
       uvloop
       uvicorn
+      websockets
     ];
+
     extra_proxy = [
       azure-identity
       azure-keyvault-secrets
       google-cloud-kms
       prisma
-      pynacl
       resend
     ];
   };
 
-  pythonImportsCheck = [ "litellm" ];
+  pythonImportsCheck = [
+    "litellm"
+    "litellm_enterprise"
+  ];
 
   # Relax dependency check on openai, may not be needed in the future
   pythonRelaxDeps = [ "openai" ];
 
   # access network
   doCheck = false;
 
+  postFixup = ''
+    # Symlink litellm_enterprise to make it discoverable
+    pushd $out/lib/python${python.pythonVersion}/site-packages
+    ln -s enterprise/litellm_enterprise litellm_enterprise
+    popd
+  '';
+
   passthru = {
     tests = { inherit (nixosTests) litellm; };
     updateScript = nix-update-script {