From 7749b8aaf4390a83068045dd98f63345ed9566d9 Mon Sep 17 00:00:00 2001 From: Lifei Zhou Date: Tue, 22 Oct 2024 14:39:58 +1100 Subject: [PATCH] fix: Simplifed langfuse auth check (#177) --- packages/exchange/.env.langfuse.local | 5 --- .../exchange/src/exchange/langfuse_wrapper.py | 37 +++++++------------ .../exchange/tests/test_langfuse_wrapper.py | 10 +++-- 3 files changed, 19 insertions(+), 33 deletions(-) diff --git a/packages/exchange/.env.langfuse.local b/packages/exchange/.env.langfuse.local index cdebcd7a16..866ab3f42e 100644 --- a/packages/exchange/.env.langfuse.local +++ b/packages/exchange/.env.langfuse.local @@ -9,8 +9,3 @@ LANGFUSE_INIT_USER_PASSWORD=localpwd LANGFUSE_INIT_ORG_ID=local-id LANGFUSE_INIT_ORG_NAME=local-org LANGFUSE_INIT_PROJECT_ID=goose - -# These variables are used by Goose -LANGFUSE_PUBLIC_KEY=publickey-local -LANGFUSE_SECRET_KEY=secretkey-local -LANGFUSE_HOST=http://localhost:3000 diff --git a/packages/exchange/src/exchange/langfuse_wrapper.py b/packages/exchange/src/exchange/langfuse_wrapper.py index 23650368b2..9788bf4640 100644 --- a/packages/exchange/src/exchange/langfuse_wrapper.py +++ b/packages/exchange/src/exchange/langfuse_wrapper.py @@ -14,29 +14,28 @@ Note: import os from typing import Callable -from dotenv import load_dotenv from langfuse.decorators import langfuse_context import sys from io import StringIO -from pathlib import Path -from functools import wraps # Add this import - - -def find_package_root(start_path: Path, marker_file: str = "pyproject.toml") -> Path: - while start_path != start_path.parent: - if (start_path / marker_file).exists(): - return start_path - start_path = start_path.parent - return None +from functools import cache, wraps + +## These are the default configurations for local Langfuse server +## Please refer to .env.langfuse.local file for local langfuse server setup configurations +DEFAULT_LOCAL_LANGFUSE_HOST = "http://localhost:3000" +DEFAULT_LOCAL_LANGFUSE_PUBLIC_KEY = "publickey-local" +DEFAULT_LOCAL_LANGFUSE_SECRET_KEY = "secretkey-local" +@cache def auth_check() -> bool: # Temporarily redirect stdout and stderr to suppress print statements from Langfuse temp_stderr = StringIO() sys.stderr = temp_stderr - # Load environment variables - load_dotenv(LANGFUSE_ENV_FILE, override=True) + # Set environment variables if not specified + os.environ.setdefault("LANGFUSE_PUBLIC_KEY", DEFAULT_LOCAL_LANGFUSE_PUBLIC_KEY) + os.environ.setdefault("LANGFUSE_SECRET_KEY", DEFAULT_LOCAL_LANGFUSE_SECRET_KEY) + os.environ.setdefault("LANGFUSE_HOST", DEFAULT_LOCAL_LANGFUSE_HOST) auth_val = langfuse_context.auth_check() @@ -45,16 +44,6 @@ def auth_check() -> bool: return auth_val -CURRENT_DIR = Path(__file__).parent -PACKAGE_ROOT = find_package_root(CURRENT_DIR) - -LANGFUSE_ENV_FILE = os.path.join(PACKAGE_ROOT, ".env.langfuse.local") if PACKAGE_ROOT else None -HAS_LANGFUSE_CREDENTIALS = False -load_dotenv(LANGFUSE_ENV_FILE, override=True) - -HAS_LANGFUSE_CREDENTIALS = auth_check() - - def observe_wrapper(*args, **kwargs) -> Callable: # noqa """ A decorator that wraps a function with Langfuse context observation if credentials are available. @@ -71,7 +60,7 @@ def observe_wrapper(*args, **kwargs) -> Callable: # noqa """ def _wrapper(fn: Callable) -> Callable: - if HAS_LANGFUSE_CREDENTIALS: + if auth_check(): @wraps(fn) def wrapped_fn(*fargs, **fkwargs): # noqa diff --git a/packages/exchange/tests/test_langfuse_wrapper.py b/packages/exchange/tests/test_langfuse_wrapper.py index 3218509785..9ac304c536 100644 --- a/packages/exchange/tests/test_langfuse_wrapper.py +++ b/packages/exchange/tests/test_langfuse_wrapper.py @@ -9,9 +9,10 @@ def mock_langfuse_context(): yield mock -@patch("exchange.langfuse_wrapper.HAS_LANGFUSE_CREDENTIALS", True) -def test_function_is_wrapped(mock_langfuse_context): +@patch("exchange.langfuse_wrapper.auth_check") +def test_function_is_wrapped(mock_auth_check, mock_langfuse_context): mock_observe = MagicMock(side_effect=lambda *args, **kwargs: lambda fn: fn) + mock_auth_check.return_value = True mock_langfuse_context.observe = mock_observe def original_function(x: int, y: int) -> int: @@ -31,9 +32,10 @@ def test_function_is_wrapped(mock_langfuse_context): mock_observe.assert_called_with("arg1", kwarg1="kwarg1") -@patch("exchange.langfuse_wrapper.HAS_LANGFUSE_CREDENTIALS", False) -def test_function_is_not_wrapped(mock_langfuse_context): +@patch("exchange.langfuse_wrapper.auth_check") +def test_function_is_not_wrapped(mock_auth_check, mock_langfuse_context): mock_observe = MagicMock(return_value=lambda f: f) + mock_auth_check.return_value = False mock_langfuse_context.observe = mock_observe @observe_wrapper("arg1", kwarg1="kwarg1")