From dcc3bbc26430198f2f0dc6f1c303e554b217d5bb Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 25 Jan 2025 17:17:32 -0800
Subject: [PATCH] (Fix) langfuse - setting `LANGFUSE_FLUSH_INTERVAL`  (#8007)

* fix langfuse flush interval

* test_get_langfuse_flush_interval

* test_get_langfuse_flush_interval
---
 litellm/integrations/langfuse/langfuse.py     | 20 ++++++++++++++--
 .../langfuse/langfuse_prompt_management.py    |  5 ++--
 .../test_langfuse_unit_tests.py               | 23 +++++++++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py
index 20d2befe65..125bf4e686 100644
--- a/litellm/integrations/langfuse/langfuse.py
+++ b/litellm/integrations/langfuse/langfuse.py
@@ -54,8 +54,8 @@ class LangFuseLogger:
             self.langfuse_host = "http://" + self.langfuse_host
         self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
         self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")
-        self.langfuse_flush_interval = (
-            os.getenv("LANGFUSE_FLUSH_INTERVAL") or flush_interval
+        self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(
+            flush_interval
         )
         http_client = _get_httpx_client()
         self.langfuse_client = http_client.client
@@ -708,6 +708,22 @@ class LangFuseLogger:
         """Check if current langfuse version supports completion start time"""
         return Version(self.langfuse_sdk_version) >= Version("2.7.3")
 
+    @staticmethod
+    def _get_langfuse_flush_interval(flush_interval: int) -> int:
+        """
+        Get the langfuse flush interval to initialize the Langfuse client
+
+        Reads `LANGFUSE_FLUSH_INTERVAL` from the environment variable.
+        If not set, uses the flush interval passed in as an argument.
+
+        Args:
+            flush_interval: The flush interval to use if LANGFUSE_FLUSH_INTERVAL is not set
+
+        Returns:
+            [int] The flush interval to use to initialize the Langfuse client
+        """
+        return int(os.getenv("LANGFUSE_FLUSH_INTERVAL") or flush_interval)
+
 
 def _add_prompt_to_generation_params(
     generation_params: dict,
diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py
index 44419896c4..faa4a63491 100644
--- a/litellm/integrations/langfuse/langfuse_prompt_management.py
+++ b/litellm/integrations/langfuse/langfuse_prompt_management.py
@@ -81,7 +81,6 @@ def langfuse_client_init(
 
     langfuse_release = os.getenv("LANGFUSE_RELEASE")
     langfuse_debug = os.getenv("LANGFUSE_DEBUG")
-    langfuse_flush_interval = os.getenv("LANGFUSE_FLUSH_INTERVAL") or flush_interval
 
     parameters = {
         "public_key": public_key,
@@ -89,7 +88,9 @@ def langfuse_client_init(
         "host": langfuse_host,
         "release": langfuse_release,
         "debug": langfuse_debug,
-        "flush_interval": langfuse_flush_interval,  # flush interval in seconds
+        "flush_interval": LangFuseLogger._get_langfuse_flush_interval(
+            flush_interval
+        ),  # flush interval in seconds
     }
 
     if Version(langfuse.version.__version__) >= Version("2.6.0"):
diff --git a/tests/logging_callback_tests/test_langfuse_unit_tests.py b/tests/logging_callback_tests/test_langfuse_unit_tests.py
index 8f0f2637c6..e9c255c1a3 100644
--- a/tests/logging_callback_tests/test_langfuse_unit_tests.py
+++ b/tests/logging_callback_tests/test_langfuse_unit_tests.py
@@ -294,6 +294,28 @@ def test_get_langfuse_tags():
     assert result == []
 
 
+
+@patch.dict(os.environ, {}, clear=True)  # Start with empty environment
+def test_get_langfuse_flush_interval():
+    """
+    Test that _get_langfuse_flush_interval correctly reads from environment variable
+    or falls back to the provided flush_interval
+    """
+    default_interval = 60
+
+    # Test when env var is not set
+    result = LangFuseLogger._get_langfuse_flush_interval(
+        flush_interval=default_interval
+    )
+    assert result == default_interval
+
+    # Test when env var is set
+    with patch.dict(os.environ, {"LANGFUSE_FLUSH_INTERVAL": "120"}):
+        result = LangFuseLogger._get_langfuse_flush_interval(
+            flush_interval=default_interval
+        )
+        assert result == 120
+
 def test_langfuse_e2e_sync(monkeypatch):
     from litellm import completion
     import litellm
@@ -320,3 +342,4 @@ def test_langfuse_e2e_sync(monkeypatch):
         time.sleep(3)
 
         assert langfuse_mock.called
+