From 7a30a4d7dafaf496de8a8cc30b054ded19279330 Mon Sep 17 00:00:00 2001
From: codgician <15964984+codgician@users.noreply.github.com>
Date: Wed, 14 May 2025 19:03:09 +0800
Subject: [PATCH] python3Packages.litellm: fix litellm_enterprise module not
 found

---
 .../python-modules/litellm/default.nix        | 24 ++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/pkgs/development/python-modules/litellm/default.nix b/pkgs/development/python-modules/litellm/default.nix
index 0be73a5e530b..5cad63057288 100644
--- a/pkgs/development/python-modules/litellm/default.nix
+++ b/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,29 +79,37 @@ 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" ];
@@ -106,6 +117,13 @@ buildPythonPackage rec {
   # 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 {