From 7b91f2182ba1b1addb97702d3b19ed74fa8c5023 Mon Sep 17 00:00:00 2001
From: Pavel Zorin <pazonec@yandex.ru>
Date: Wed, 19 Nov 2025 19:23:37 +0100
Subject: [PATCH] from https://github.com/BerriAI/litellm/issues/16810 
 LITELLM_LOCAL_MODEL_COST_MAP: True

---
 .github/workflows/e2e_tests.yml          | 4 ++++
 cognee/tests/test_cognee_server_start.py | 4 +---
 2 files changed, 5 insertions(+), 3 deletions(-)

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):