diff --git a/.github/workflows/e2e_tests.yml b/.github/workflows/e2e_tests.yml
index 3dea2548c..747be2ecf 100644
--- a/.github/workflows/e2e_tests.yml
+++ b/.github/workflows/e2e_tests.yml
@@ -374,6 +374,8 @@ jobs:
       - name: Run Entity Extraction Test
         env:
           ENV: 'dev'
+          # remove when https://github.com/BerriAI/litellm/issues/16810 is fixed
+          LITELLM_LOCAL_MODEL_COST_MAP: "True" 
           LLM_MODEL: ${{ secrets.LLM_MODEL }}
           LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
           LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
@@ -402,6 +404,8 @@ jobs:
       - name: Run Feedback Enrichment Test
         env:
           ENV: 'dev'
+          # remove when https://github.com/BerriAI/litellm/issues/16810 is fixed
+          LITELLM_LOCAL_MODEL_COST_MAP: "True" 
           LLM_MODEL: ${{ secrets.LLM_MODEL }}
           LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
           LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
diff --git a/cognee/tests/test_cognee_server_start.py b/cognee/tests/test_cognee_server_start.py
index 2ae444e6b..4ec5a254d 100644
--- a/cognee/tests/test_cognee_server_start.py
+++ b/cognee/tests/test_cognee_server_start.py
@@ -9,9 +9,7 @@ from pathlib import Path
 import sys
 import uuid
 
-from tenacity import retry, stop_after_attempt, wait_fixed, wait_incrementing
-
-from cognee.api.health import HealthStatus
+from tenacity import retry, stop_after_attempt, wait_incrementing
 
 
 class TestCogneeServerStart(unittest.TestCase):