from https://github.com/BerriAI/litellm/issues/16810 LITELLM_LOCAL_MODEL_COST_MAP: True

.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 }}

.gitignore

@@ -13,6 +13,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+mise.toml
+
 full_run.ipynb
 
 # C extensions

cognee/tests/test_cognee_server_start.py

@@ -1,3 +1,4 @@
+from logging import getLogger
 import unittest
 import subprocess
 import time
@@ -8,6 +9,8 @@ from pathlib import Path
 import sys
 import uuid
 
+from tenacity import retry, stop_after_attempt, wait_incrementing
+
 
 class TestCogneeServerStart(unittest.TestCase):
     @classmethod
@@ -27,9 +30,7 @@ class TestCogneeServerStart(unittest.TestCase):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             preexec_fn=os.setsid,
-        )
-        # Give the server some time to start
-        time.sleep(35)
+        )                
 
         # Check if server started with errors
         if cls.server_process.poll() is not None:
@@ -49,12 +50,22 @@ class TestCogneeServerStart(unittest.TestCase):
                 cls.server_process.terminate()
             cls.server_process.wait()
 
+    
+    @retry(wait=wait_incrementing(start=5, increment=3), stop=stop_after_attempt(5))
+    def wait_for_server_startup(self):
+        print("Waiting for server to start...")
+        response = requests.get("http://localhost:8000/health", timeout=15)
+        response.raise_for_status()
+        
+        status_str = response.json().get("status")
+        if response.status_code == 200 and status_str == "ready":
+            return True
+        else:
+            getLogger().info(f"Health check status {response.status_code} : {status_str}, retrying...")
+        raise requests.exceptions.RequestException(f"Server is {status_str}")
+        
     def test_server_is_running(self):
-        """Test that the server is running and can accept connections."""
-        # Test health endpoint
-        health_response = requests.get("http://localhost:8000/health", timeout=15)
-        self.assertIn(health_response.status_code, [200])
-
+        self.wait_for_server_startup()
         # Test root endpoint
         root_response = requests.get("http://localhost:8000/", timeout=15)
         self.assertEqual(root_response.status_code, 200)