Ap 693: Add the updated Callbackhandler (#67)

Author: yzhoubk

mokelumne/util/image_describer.py

@@ -7,7 +7,7 @@
 
 from botocore.exceptions import ClientError
 from langchain_core.messages import HumanMessage, SystemMessage
-from langfuse.langchain import CallbackHandler
+from mokelumne.util import langfuse
 
 
 logger = logging.getLogger(__name__)
@@ -17,11 +17,13 @@
 class ImageDescriber:
     """Generates a description of a given image using a large language model."""
 
-    def __init__(self, model, prompt: str):
+    def __init__(self, model, prompt: str, conn_id: str = 'langfuse_default'):
         """Initialise an image describer with a given model and prompt."""
+
         self.model = model
         self.prompt = prompt
-        self.langfuse_handler = CallbackHandler()
+        self.conn_id = conn_id
+        self.langfuse_handler = langfuse.get_langfuse_callback_handler(self.conn_id)
         self.sys_msg = SystemMessage(prompt)
 
     def describe(self, record: dict[str, str]) -> dict[str, str]:

mokelumne/util/langfuse.py

@@ -8,6 +8,7 @@
 
 from airflow.sdk import BaseHook
 from langfuse import Langfuse
+from langfuse.langchain import CallbackHandler
 
 Prompt = namedtuple('Prompt', ['prompt', 'version'])
 Prompt.__doc__ += ': Contains a LLM prompt for generating image descriptions.'
@@ -31,6 +32,12 @@ def _get_langfuse_connection_settings(conn_id: str) -> tuple[str, str, str]:
         )
     return base_url, public_key, secret_key
 
+def get_langfuse_callback_handler(conn_id: str) -> CallbackHandler:
+    """Return a LangChain CallbackHandler configured from the given Airflow connection."""
+    get_langfuse_client(conn_id)
+    _, public_key, _ = _get_langfuse_connection_settings(conn_id)
+    return CallbackHandler(public_key=public_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)

test/unit/test_langfuse.py

@@ -124,6 +124,26 @@ def test_passes_release_and_environment(self, monkeypatch):
         assert 'release' in call_kwargs
 
 
+@pytest.mark.usefixtures("mock_conn_settings")
+class TestGetLangfuseCallbackHandler:
+    """Tests for get_langfuse_callback_handler."""
+
+    def test_callback_handler_uses_given_conn_id(self, monkeypatch):
+        """Ensure callback handler setup forwards custom conn_id."""
+        callback_handler_cls = Mock(return_value=Mock())
+        get_client = Mock(return_value=Mock())
+        get_conn_settings = Mock(return_value=FAKE_CONN_SETTINGS)
+
+        monkeypatch.setattr(langfuse, 'CallbackHandler', callback_handler_cls)
+        monkeypatch.setattr(langfuse, 'get_langfuse_client', get_client)
+        monkeypatch.setattr(langfuse, '_get_langfuse_connection_settings', get_conn_settings)
+
+        langfuse.get_langfuse_callback_handler('langfuse_custom')
+
+        get_client.assert_called_once_with('langfuse_custom')
+        get_conn_settings.assert_called_once_with('langfuse_custom')
+
+
 @pytest.mark.usefixtures("mock_conn_settings")
 class TestGetPrompt:
     """Tests for get_prompt."""