fix: langfuse callback

Signed-off-by: thiswillbeyourgithub <[email protected]>

Author: thiswillbeyourgithub

wdoc/utils/llm.py

@@ -22,34 +22,10 @@
 from .typechecker import optional_typecheck
 from .flags import is_verbose, is_private
 from .env import WDOC_PRIVATE_MODE
-from .misc import version_holder
+from .misc import langfuse_callback_holder
 
 TESTING_LLM = "testing/testing"
 
-if (
-        "LANGFUSE_PUBLIC_KEY" in os.environ and
-        "LANGFUSE_SECRET_KEY" in os.environ and
-        "LANGFUSE_HOST" in os.environ
-) and not is_private:
-    red("Activating langfuse callbacks")
-    try:
-        # # litellm's callbacks seem more flawed than langchain's
-        # import langfuse
-        # litellm.success_callback = ["langfuse"]
-        # litellm.failure_callback = ["langfuse"]
-
-        from langfuse.callback import CallbackHandler as LangfuseCallback
-        langfuse_handler = [LangfuseCallback(
-            secret_key=os.environ["LANGFUSE_SECRET_KEY"],
-            public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
-            host=os.environ["LANGFUSE_HOST"],
-            session_id=str(uuid.uuid4()),
-            version=version_holder[0],
-        )]
-    except Exception as e:
-        red(f"Failed to setup langfuse callback, make sure package 'langfuse' is installed. The error was: ''{e}'")
-        langfuse_handler = []
-
 
 @optional_typecheck
 def load_llm(
@@ -130,7 +106,7 @@ def load_llm(
         assert os.environ[f"{backend.upper()}_API_KEY"] == "REDACTED_BECAUSE_WDOC_IN_PRIVATE_MODE"
 
         assert os.environ[f"{backend.upper()}_API_KEY"] == "REDACTED_BECAUSE_WDOC_IN_PRIVATE_MODE"
-        assert not langfuse_handler, "Private argument but langfuse_handler appears set. Something went wrong so crashing just to be safe."
+        assert not langfuse_callback_holder, "Private argument but langfuse_handler appears set. Something went wrong so crashing just to be safe."
     else:
         assert not WDOC_PRIVATE_MODE
         assert "WDOC_PRIVATE_MODE" not in os.environ or os.environ["WDOC_PRIVATE_MODE"] == "false"
@@ -146,7 +122,7 @@ def load_llm(
             cache=llm_cache,
             disable_streaming=True,  # Not needed and might break cache
             verbose=llm_verbosity,
-            callbacks=[PriceCountingCallback(verbose=llm_verbosity)] + langfuse_handler,
+            callbacks=[PriceCountingCallback(verbose=llm_verbosity)] + langfuse_callback_holder,
             **extra_model_args,
         )
     else:
@@ -160,7 +136,7 @@ def load_llm(
             cache=llm_cache,
             verbose=llm_verbosity,
             tags=tags,
-            callbacks=[PriceCountingCallback(verbose=llm_verbosity)] + langfuse_handler,
+            callbacks=[PriceCountingCallback(verbose=llm_verbosity)] + langfuse_callback_holder,
             **extra_model_args,
         )
         litellm.drop_params = True

wdoc/utils/misc.py

@@ -6,6 +6,7 @@
 from typing import List, Union, Callable, get_type_hints, Literal
 from joblib import Memory
 from joblib import hash as jhash
+import uuid
 import socket
 import os
 import json
@@ -29,7 +30,7 @@
 
 from .logger import whi, red, yel, cache_dir
 from .typechecker import optional_typecheck
-from .flags import is_verbose, is_debug
+from .flags import is_verbose, is_debug, is_private
 from .errors import UnexpectedDocDictArgument
 from .env import WDOC_NO_MODELNAME_MATCHING, WDOC_STRICT_DOCDICT, WDOC_EXPIRE_CACHE_DAYS, WDOC_IMPORT_TYPE
 
@@ -759,4 +760,29 @@ def thinking_answer_parser(output: str, strict: bool = False) -> dict:
             }
 
 # this will contain wdoc's version to be used by langfuse's callback without circular imports
-version_holder = []
+langfuse_callback_holder = []
+def create_langfuse_callback(version: str) -> None:
+    if (
+            "LANGFUSE_PUBLIC_KEY" in os.environ and
+            "LANGFUSE_SECRET_KEY" in os.environ and
+            "LANGFUSE_HOST" in os.environ
+    ) and not is_private:
+        red("Activating langfuse callbacks")
+        try:
+            # # litellm's callbacks seem more flawed than langchain's
+            # import langfuse
+            # litellm.success_callback = ["langfuse"]
+            # litellm.failure_callback = ["langfuse"]
+
+            from langfuse.callback import CallbackHandler as LangfuseCallback
+            langfuse_callback = LangfuseCallback(
+                secret_key=os.environ["LANGFUSE_SECRET_KEY"],
+                public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
+                host=os.environ["LANGFUSE_HOST"],
+                session_id=str(uuid.uuid4()),
+                version=version,
+            )
+            langfuse_callback_holder.append(langfuse_callback)
+        except Exception as e:
+            red(f"Failed to setup langfuse callback, make sure package 'langfuse' is installed. The error was: ''{e}'")
+

wdoc/wdoc.py

@@ -32,7 +32,7 @@
     extra_args_types, disable_internet,
     set_func_signature, query_eval_cache,
     thinking_answer_parser, DocDict,
-    version_holder,
+    create_langfuse_callback,
 )
 from .utils.prompts import prompts
 from .utils.tasks.query import refilter_docs, check_intermediate_answer, parse_eval_output, pbar_chain, pbar_closer, collate_intermediate_answers, semantic_batching, sieve_documents
@@ -1973,4 +1973,4 @@ def p(message: str) -> None:
     sys.excepthook = handle_exception
     faulthandler.enable()
 
-version_holder.append(wdoc.VERSION)
+create_langfuse_callback(wdoc.VERSION)