Ap 681 : Create Airflow connection for Langfuse (#64)

Author: yzhoubk

README.md

@@ -86,6 +86,7 @@ Important environment variables for our build/environment:
 | `AIRFLOW__API__BASE_URL` | Where Airflow's webserver is reachable. | `AIRFLOW__API__BASE_URL="http://localhost:8080"` |
 | `AIRFLOW__CORE__FERNET_KEY` | [Fernet](https://airflow.apache.org/docs/apache-airflow/stable/security/secrets/fernet.html) encryption key used to encrypt Airflow secrets | `AIRFLOW__CORE__FERNET_KEY="somebase64value="` |
 | `AIRFLOW__API_AUTH__JWT_SECRET` | Secret key used to sign JWT tokens for Airflow's API authentication. The default value used in development and testing should be replaced in production. | `AIRFLOW__API_AUTH__JWT_SECRET="some32bytesecret"` |
+| `AIRFLOW_CONN_LANGFUSE_DEFAULT` | Airflow connection string for Langfuse access.<br>Note: This will create a `langfuse_default` conn_id, will not store the connection in the Airflow metastore, and will override any other connection settings.  | `AIRFLOW_CONN_LANGFUSE_DEFAULT='{"conn_type":"langfuse","host":"us.cloud.langfuse.com","schema":"https","login":"pk-lf-blah-blah-blah","password":"ssk-lf-blah-blah-blah"}'`|
 | `AIRFLOW_CONN_TIND_DEFAULT` | (Optional override) Airflow connection json string for TIND access.<br>Note: This will create a `tind_default` conn_id, will not store the connection in the Airflow metastore, and will override any other connection settings.  | `AIRFLOW_CONN_TIND_DEFAULT='{"conn_type": "http","password": "your-tind-key-here","host": "https://digicoll.lib.berkeley.edu/api/v1","schema": "https"}'` |
 | `OIDC_CLIENT_SECRET` | Client secret for OIDC authentication.  Used by the Airflow webserver to authenticate OIDC token requests.  In development, also used by `keycloak-config-cli` to configure the client secret.  This should match Keycloak configuration in development and testing, and CalNet in production. | `OIDC_CLIENT_SECRET="some32charactersecret"` |
 | `OIDC_NAME` | Name appended to the OIDC login button | `OIDC_NAME="keycloak"` |
@@ -98,9 +99,6 @@ Important environment variables for our build/environment:
 | `TIND_API_URL` | URL for TIND access | `TIND_API_URL="https://digicoll.lib.berkeley.edu/api/v1"` |
 | `TIND_IIIF_MANIFEST_URL_PATTERN` | URL pattern for TIND IIIF manifests | `TIND_IIIF_MANIFEST_URL_PATTERN="https://digicoll.lib.berkeley.edu/record/{tind_id}/export/iiif_manifest"` |
 | `MOKELUMNE_TIND_DOWNLOAD_DIR` | Path for downloaded image cache | `MOKELUMNE_TIND_DOWNLOAD_DIR="/some/path/to/download/to"` |
-|`LANGFUSE_HOST`|Host for Langfuse|`LANGFUSE_HOST="https://us.cloud.langfuse.com"`|
-|`LANGFUSE_SECRET_KEY`|sets langfuse secret key|`LANGFUSE_SECRET_KEY="sk-lf-blah-blah-blah"`|
-|`LANGFUSE_PUBLIC_KEY`|sets langfuse public key|`LANGFUSE_PUBLIC_KEY="pk-lf-blah-blah-blah"`|
 |`AWS_ENDPOINT_URL`|AWS endpoint (don't forget the `https://`!)|`AWS_ENDPOINT_URL="https://bedrock-runtime.us-west-1.amazonaws.com"`|
 |`AWS_DEFAULT_REGION`|The AWS region to use; you probably want `us-west-1`.|`AWS_DEFAULT_REGION=us-west-1`|
 |`AWS_BEARER_TOKEN_BEDROCK`|The IAM credential to use to access AWS. Use a short-term API key.<br>The key will expire after AWS console logout or 12 hours (whichever comes first).<br>Make sure that your region for the key matches the region above.|`AWS_BEARER_TOKEN_BEDROCK="bedrock-api-key-blah-blah-blah"`|

example.env

@@ -5,6 +5,8 @@ AIRFLOW__API_AUTH__JWT_SECRET=
 # @see https://airflow.apache.org/docs/apache-airflow/stable/security/secrets/fernet.html#generating-fernet-key
 AIRFLOW__CORE__FERNET_KEY=
 
+AIRFLOW_CONN_LANGFUSE_DEFAULT='{"conn_type":"langfuse","host":"us.cloud.langfuse.com","schema":"https","login":"pk-lf-blah-blah-blah","password":"ssk-lf-blah-blah-blah"}'
+
 # @see README.md for how to use the optional AIRFLOW_CONN override ariable.
 # AIRFLOW_CONN_TIND_DEFAULT='{"conn_type": "http","password": "your-tind-key-here","host": "https://digicoll.lib.berkeley.edu/api/v1","schema": "https"}'
 
@@ -29,10 +31,6 @@ TIND_API_KEY=your-tind-key-here
 TIND_API_URL=https://digicoll.lib.berkeley.edu/api/v1
 TIND_IIIF_MANIFEST_URL_PATTERN=https://digicoll.lib.berkeley.edu/record/{tind_id}/export/iiif_manifest
 
-LANGFUSE_HOST=https://us.cloud.langfuse.com
-LANGFUSE_SECRET_KEY=sk-lf-blah-blah-blah
-LANGFUSE_PUBLIC_KEY=pk-lf-blah-blah-blah
-
 AWS_ENDPOINT_URL=https://bedrock-runtime.us-west-1.amazonaws.com
 AWS_DEFAULT_REGION=us-west-1
 AWS_BEARER_TOKEN_BEDROCK=bedrock-api-key-blah-blah-blah

mokelumne/dags/gen_llm_image_descriptions.py

@@ -262,7 +262,7 @@ def get_prompt() -> dict[str, str]:
             context = get_current_context()
             prompt = langfuse.get_prompt(
                 name=context["params"]["langfuse_prompt_name"],
-                version_or_label=context["params"]["langfuse_prompt_version_or_label"],
+                version_or_label=context["params"]["langfuse_prompt_version_or_label"]
             )
 
             return {"prompt": prompt.prompt, "version": prompt.version}

mokelumne/util/langfuse.py

@@ -6,9 +6,9 @@
 from collections import namedtuple
 from os import environ as ENV
 
+from airflow.sdk import BaseHook
 from langfuse import Langfuse
 
-
 Prompt = namedtuple('Prompt', ['prompt', 'version'])
 Prompt.__doc__ += ': Contains a LLM prompt for generating image descriptions.'
 Prompt.prompt.__doc__ = 'The LLM prompt.'
@@ -17,10 +17,34 @@
 logger = logging.getLogger(__name__)
 
 
-def get_prompt(name: str, version_or_label: int | str) -> Prompt:
+def _get_langfuse_connection_settings(conn_id: str) -> tuple[str, str, str]:
+    """Return host/public/secret key tuple from the Langfuse Airflow connection."""
+    conn = BaseHook.get_connection(conn_id)
+    base_url = f'{conn.schema}://{conn.host}'
+    public_key = conn.login
+    secret_key = conn.password
+
+    if not public_key or not secret_key:
+        raise ValueError(
+            f'Missing Langfuse credentials in Airflow connection {conn_id}. '
+            'Set login/password on the connection.'
+        )
+    return base_url, public_key, secret_key
+
+def get_langfuse_client(conn_id: str) -> Langfuse:
+    """Return a Langfuse client configured from the ``langfuse_default`` Airflow connection."""
+    base_url, public_key, secret_key = _get_langfuse_connection_settings(conn_id)
+    return Langfuse(
+        base_url=base_url,
+        public_key=public_key,
+        secret_key=secret_key,
+        release=importlib.metadata.version('mokelumne'),
+        environment=ENV.get('DEPLOYMENT_ID', 'default'),
+    )
+
+def get_prompt(name: str, version_or_label: int | str, conn_id: str = 'langfuse_default') -> Prompt:
     """Return the current prompt to use."""
-    langfuse = Langfuse(release=importlib.metadata.version('mokelumne'),
-                        environment=ENV.get('DEPLOYMENT_ID', 'default'))
+    langfuse = get_langfuse_client(conn_id)
     if isinstance(version_or_label, int):
         logger.debug(
             f"Getting Langfuse prompt {name}, version {version_or_label}"

test/unit/test_langfuse.py

@@ -2,17 +2,19 @@
 
 from unittest.mock import Mock
 import pytest
-
 from mokelumne.util import langfuse
 
-
 DEFAULT_TEST_PROMPT: str = "A test prompt."
 """The prompt used if none is specified to the mocked Langfuse object."""
 
-
 DEFAULT_TEST_VERSION: int = 67
 """The version used if none is specified to the mocked Langfuse object."""
 
+FAKE_BASE_URL: str = "https://langfuse.example.com"
+FAKE_PUBLIC_KEY: str = "pk-test-123"
+FAKE_SECRET_KEY: str = "sk-test-456"
+FAKE_CONN_SETTINGS: tuple[str, str, str] = (FAKE_BASE_URL, FAKE_PUBLIC_KEY, FAKE_SECRET_KEY)
+
 
 class MockResult:
     """An object for return by our mock Langfuse."""
@@ -37,14 +39,100 @@ def __init__(self, **kwargs):
         self.get_prompt = Mock(return_value=mocked_result)
 
 
-class TestLangfuse:
-    """Tests for the Mokelumne Langfuse integration utility class."""
+@pytest.fixture
+def mock_conn_settings(monkeypatch):
+    """Patch _get_langfuse_connection_settings to avoid hitting Airflow."""
+    monkeypatch.setattr(
+        langfuse, '_get_langfuse_connection_settings', lambda conn_id: FAKE_CONN_SETTINGS
+    )
+
+
+class MockConnection:
+    """Mock Airflow connection object."""
+    def __init__(self, host, schema=None, login=None, password=None):
+        self.host = host
+        self.schema = schema
+        self.login = login
+        self.password = password
+
+
+class TestGetLangfuseConnectionSettings:
+    """Tests for _get_langfuse_connection_settings."""
+
+    def test_extracts_credentials_from_airflow_connection(self, monkeypatch):
+        """Ensure Langfuse credentials are extracted from Airflow connection fields."""
+        mock_conn = MockConnection(
+            host="langfuse.example.com",
+            schema="https",
+            login="pk-123",
+            password="sk-456"
+        )
+        monkeypatch.setattr(
+            langfuse.BaseHook, 'get_connection', Mock(return_value=mock_conn)
+        )
+
+        host, public_key, secret_key = langfuse._get_langfuse_connection_settings('langfuse_default')
+
+        assert host == "https://langfuse.example.com"
+        assert public_key == "pk-123"
+        assert secret_key == "sk-456"
+
+    def test_raises_on_missing_login_credentials(self, monkeypatch):
+        """Ensure ValueError is raised when login (public key) is missing."""
+        mock_conn = MockConnection(
+            host="langfuse.example.com",
+            schema="https",
+            login=None,
+            password="sk-456"
+        )
+        monkeypatch.setattr(
+            langfuse.BaseHook, 'get_connection', Mock(return_value=mock_conn)
+        )
+
+        with pytest.raises(ValueError, match="Missing Langfuse credentials"):
+            langfuse._get_langfuse_connection_settings('langfuse_default')
+
+
+@pytest.mark.usefixtures("mock_conn_settings")
+class TestGetLangfuseClient:
+    """Tests for get_langfuse_client."""
+
+    def test_returns_langfuse_instance(self, monkeypatch):
+        """Ensure get_langfuse_client constructs a Langfuse with the connection settings."""
+        mock_langfuse_cls = Mock(return_value=Mock())
+        monkeypatch.setattr(langfuse, 'Langfuse', mock_langfuse_cls)
+
+        client = langfuse.get_langfuse_client('langfuse_default')
+
+        mock_langfuse_cls.assert_called_once()
+        call_kwargs = mock_langfuse_cls.call_args.kwargs
+        assert call_kwargs['base_url'] == FAKE_BASE_URL
+        assert call_kwargs['public_key'] == FAKE_PUBLIC_KEY
+        assert call_kwargs['secret_key'] == FAKE_SECRET_KEY
+        assert client is mock_langfuse_cls.return_value
+
+    def test_passes_release_and_environment(self, monkeypatch):
+        """Ensure release and environment are forwarded to the Langfuse constructor."""
+        mock_langfuse_cls = Mock(return_value=Mock())
+        monkeypatch.setattr(langfuse, 'Langfuse', mock_langfuse_cls)
+        monkeypatch.setenv('DEPLOYMENT_ID', 'staging')
+
+        langfuse.get_langfuse_client('langfuse_default')
+
+        call_kwargs = mock_langfuse_cls.call_args.kwargs
+        assert call_kwargs['environment'] == 'staging'
+        assert 'release' in call_kwargs
+
+
+@pytest.mark.usefixtures("mock_conn_settings")
+class TestGetPrompt:
+    """Tests for get_prompt."""
 
     def test_simple_prompt(self, monkeypatch):
         """Test a simple call of the `get_prompt` method."""
         with monkeypatch.context() as m:
             m.setattr(langfuse, 'Langfuse', MockLangfuse)
-            result = langfuse.get_prompt('test', 'production')
+            result = langfuse.get_prompt('test', 'production', 'langfuse_default')
         assert result.prompt == DEFAULT_TEST_PROMPT
         assert result.version == DEFAULT_TEST_VERSION
 
@@ -58,7 +146,7 @@ def test_label_is_used(self, monkeypatch):
 
         with monkeypatch.context() as m:
             m.setattr(langfuse, 'Langfuse', mock)
-            langfuse.get_prompt(name, label)
+            langfuse.get_prompt(name, label, 'langfuse_default')
 
         lf_object.get_prompt.assert_called_with(name, label=label)
 
@@ -72,6 +160,6 @@ def test_version_is_used(self, monkeypatch):
 
         with monkeypatch.context() as m:
             m.setattr(langfuse, 'Langfuse', mock)
-            langfuse.get_prompt(name, version)
+            langfuse.get_prompt(name, version, 'langfuse_default')
 
         lf_object.get_prompt.assert_called_with(name, version=version)