feat(integrations): add auth check to openai and langchain (#199)

Author: maxdeichmann

langfuse/callback.py

@@ -51,10 +51,12 @@ def __init__(
             self.trace = statefulClient
             self.runs = {}
             self.rootSpan = None
+            self.langfuse = None
 
         elif statefulClient and isinstance(statefulClient, StatefulSpanClient):
             self.runs = {}
             self.rootSpan = statefulClient
+            self.langfuse = None
             self.trace = StatefulTraceClient(
                 statefulClient.client,
                 statefulClient.trace_id,
@@ -98,6 +100,22 @@ def flush(self):
         else:
             self.log.debug("There was no trace yet, hence no flushing possible.")
 
+    def auth_check(self):
+        if self.langfuse is not None:
+            return self.langfuse.auth_check()
+        elif self.trace is not None:
+            projects = self.trace.client.projects.get()
+            if len(projects.data) == 0:
+                raise Exception("No projects found for the keys.")
+            return True
+        elif self.rootSpan is not None:
+            projects = self.rootSpan.client.projects.get()
+            if len(projects) == 0:
+                raise Exception("No projects found for the keys.")
+            return True
+
+        return False
+
     def setNextSpan(self, id: str):
         self.nextSpanId = id
 

langfuse/openai.py

@@ -346,3 +346,10 @@ def register_tracing(self):
 
 modifier = OpenAILangfuse()
 modifier.register_tracing()
+
+
+def auth_check():
+    if modifier._langfuse is None:
+        modifier.initialize()
+
+    return modifier._langfuse.auth_check()

tests/test_openai.py

@@ -2,7 +2,7 @@
 import pytest
 from langfuse.client import Langfuse
 from langfuse.model import CreateTrace
-from langfuse.openai import _is_openai_v1, _is_streaming_response, openai, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI
+from langfuse.openai import _is_openai_v1, _is_streaming_response, openai, AsyncOpenAI, AzureOpenAI, AsyncAzureOpenAI, auth_check
 from openai import APIConnectionError
 
 from tests.utils import create_uuid, get_api

tests/test_sdk_setup.py

@@ -10,8 +10,9 @@
 from langfuse.callback import CallbackHandler
 
 from langfuse.client import Langfuse
+from langfuse.model import CreateTrace
 from tests.test_task_manager import get_host
-from langfuse.openai import _is_openai_v1, openai
+from langfuse.openai import _is_openai_v1, auth_check, openai
 
 chat_func = openai.chat.completions.create if _is_openai_v1() else openai.ChatCompletion.create
 
@@ -226,6 +227,27 @@ def test_openai_default():
     os.environ["LANGFUSE_HOST"] = host
 
 
+def test_openai_auth_check():
+    assert auth_check() is True
+
+
+def test_openai_auth_check_failing_key():
+    secret_key = os.environ["LANGFUSE_SECRET_KEY"]
+    os.environ.pop("LANGFUSE_SECRET_KEY")
+
+    importlib.reload(langfuse)
+    importlib.reload(langfuse.openai)
+
+    from langfuse.openai import openai
+
+    openai.langfuse_secret_key = "test"
+
+    with pytest.raises(UnauthorizedError):
+        auth_check()
+
+    os.environ["LANGFUSE_SECRET_KEY"] = secret_key
+
+
 def test_openai_configured(httpserver: HTTPServer):
     httpserver.expect_request("/api/public/ingestion", method="POST").respond_with_response(Response(status=200))
     host = get_host(httpserver.url_for("/api/public/ingestion"))
@@ -270,6 +292,7 @@ def test_openai_configured(httpserver: HTTPServer):
 
 def test_client_init_workers_5():
     langfuse = Langfuse(threads=5)
+    langfuse.flush()
 
     assert langfuse.task_manager._threads == 5
 
@@ -287,16 +310,54 @@ def test_auth_check():
 
     assert langfuse.auth_check() is True
 
+    langfuse.flush()
+
 
 def test_wrong_key_auth_check():
     langfuse = Langfuse(debug=False, secret_key="test")
 
     with pytest.raises(UnauthorizedError):
         langfuse.auth_check()
 
+    langfuse.flush()
+
+
+def test_auth_check_callback():
+    langfuse = CallbackHandler(debug=False)
+
+    assert langfuse.auth_check() is True
+    langfuse.flush()
+
+
+def test_auth_check_callback_stateful():
+    langfuse = Langfuse(debug=False)
+    trace = langfuse.trace(CreateTrace(name="name"))
+    handler = trace.get_langchain_handler()
+
+    assert handler.auth_check() is True
+    handler.flush()
+
+
+def test_wrong_key_auth_check_callback():
+    langfuse = CallbackHandler(debug=False, secret_key="test")
+
+    with pytest.raises(UnauthorizedError):
+        langfuse.auth_check()
+    langfuse.flush()
+
 
 def test_wrong_url_auth_check():
     langfuse = Langfuse(debug=False, host="http://localhost:4000/")
 
     with pytest.raises(httpx.ConnectError):
         langfuse.auth_check()
+
+    langfuse.flush()
+
+
+def test_wrong_url_auth_check_callback():
+    langfuse = CallbackHandler(debug=False, host="http://localhost:4000/")
+
+    with pytest.raises(httpx.ConnectError):
+        langfuse.auth_check()
+    langfuse.flush()