refactor: move langfuse wrapper to a module in exchange instead of a package (#138)

Co-authored-by: Alice Hau <ahau@squareup.com>

packages/exchange/.env.langfuse.local

@@ -0,0 +1,16 @@
+# These variables are default initialization variables for locally hosted Langfuse server
+LANGFUSE_INIT_PROJECT_NAME=goose-local
+LANGFUSE_INIT_PROJECT_PUBLIC_KEY=publickey-local
+LANGFUSE_INIT_PROJECT_SECRET_KEY=secretkey-local
+LANGFUSE_INIT_USER_EMAIL=local@block.xyz
+LANGFUSE_INIT_USER_NAME=localdev
+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

packages/exchange/pyproject.toml

@@ -13,6 +13,8 @@ dependencies = [
     "tiktoken>=0.7.0",
     "httpx>=0.27.0",
     "tenacity>=9.0.0",
+    "python-dotenv>=1.0.1",
+    "langfuse>=2.38.2"
 ]
 
 [tool.hatch.build.targets.wheel]

packages/exchange/src/exchange/exchange.py

@@ -3,6 +3,7 @@ import traceback
 from copy import deepcopy
 from typing import Mapping
 from attrs import define, evolve, field, Factory
+from exchange.langfuse_wrapper import observe_wrapper
 from tiktoken import get_encoding
 
 from exchange.checkpoint import Checkpoint, CheckpointData
@@ -126,6 +127,7 @@ class Exchange:
 
         return response
 
+    @observe_wrapper()
     def call_function(self, tool_use: ToolUse) -> ToolResult:
         """Call the function indicated by the tool use"""
         tool = self._toolmap.get(tool_use.name)

packages/exchange/src/exchange/langfuse_wrapper.py

@@ -0,0 +1,84 @@
+"""
+Langfuse Integration Module
+
+This module provides integration with Langfuse, a tool for monitoring and tracing LLM applications.
+
+Usage:
+    Import this module to enable Langfuse integration.
+    It automatically checks for Langfuse credentials in the .env.langfuse file and for a running Langfuse server.
+    If these are found, it will set up the necessary client and context for tracing.
+
+Note:
+    Run setup_langfuse.sh which automates the steps for running local Langfuse.
+"""
+
+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
+
+
+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)
+
+    auth_val = langfuse_context.auth_check()
+
+    # Restore stderr
+    sys.stderr = sys.__stderr__
+    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")
+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.
+
+    If Langfuse credentials were found, the function will be wrapped with Langfuse's observe method.
+    Otherwise, the function will be returned as-is.
+
+    Args:
+        *args: Positional arguments to pass to langfuse_context.observe.
+        **kwargs: Keyword arguments to pass to langfuse_context.observe.
+
+    Returns:
+        Callable: The wrapped function if credentials are available, otherwise the original function.
+    """
+
+    def _wrapper(fn: Callable) -> Callable:
+        if HAS_LANGFUSE_CREDENTIALS:
+
+            @wraps(fn)
+            def wrapped_fn(*fargs, **fkwargs):  # noqa
+                return langfuse_context.observe(*args, **kwargs)(fn)(*fargs, **fkwargs)
+
+            return wrapped_fn
+        else:
+            return fn
+
+    return _wrapper

packages/exchange/src/exchange/providers/anthropic.py

@@ -7,6 +7,7 @@ from exchange.content import Text, ToolResult, ToolUse
 from exchange.providers.base import Provider, Usage
 from tenacity import retry, wait_fixed, stop_after_attempt
 from exchange.providers.utils import retry_if_status, raise_for_status
+from exchange.langfuse_wrapper import observe_wrapper
 
 ANTHROPIC_HOST = "https://api.anthropic.com/v1/messages"
 
@@ -122,6 +123,7 @@ class AnthropicProvider(Provider):
             messages_spec.append(converted)
         return messages_spec
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/src/exchange/providers/bedrock.py

@@ -15,6 +15,7 @@ from exchange.providers import Provider, Usage
 from tenacity import retry, wait_fixed, stop_after_attempt
 from exchange.providers.utils import raise_for_status, retry_if_status
 from exchange.tool import Tool
+from exchange.langfuse_wrapper import observe_wrapper
 
 SERVICE = "bedrock-runtime"
 UTC = timezone.utc
@@ -175,6 +176,7 @@ class BedrockProvider(Provider):
         )
         return cls(client=client)
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/src/exchange/providers/databricks.py

@@ -11,7 +11,7 @@ from exchange.providers.utils import (
     tools_to_openai_spec,
 )
 from exchange.tool import Tool
-
+from exchange.langfuse_wrapper import observe_wrapper
 
 retry_procedure = retry(
     wait=wait_fixed(2),
@@ -67,6 +67,7 @@ class DatabricksProvider(Provider):
             total_tokens=total_tokens,
         )
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/src/exchange/providers/google.py

@@ -7,6 +7,8 @@ from exchange.content import Text, ToolResult, ToolUse
 from exchange.providers.base import Provider, Usage
 from tenacity import retry, wait_fixed, stop_after_attempt
 from exchange.providers.utils import raise_for_status, retry_if_status, encode_image
+from exchange.langfuse_wrapper import observe_wrapper
+
 
 GOOGLE_HOST = "https://generativelanguage.googleapis.com/v1beta"
 
@@ -131,6 +133,7 @@ class GoogleProvider(Provider):
 
         return messages_spec
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/src/exchange/providers/groq.py

@@ -1,5 +1,6 @@
 import os
 
+from exchange.langfuse_wrapper import observe_wrapper
 import httpx
 
 from exchange.message import Message
@@ -64,6 +65,7 @@ class GroqProvider(Provider):
             total_tokens=total_tokens,
         )
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/src/exchange/providers/openai.py

@@ -14,6 +14,7 @@ from exchange.providers.utils import (
 from exchange.tool import Tool
 from tenacity import retry, wait_fixed, stop_after_attempt
 from exchange.providers.utils import retry_if_status
+from exchange.langfuse_wrapper import observe_wrapper
 
 OPENAI_HOST = "https://api.openai.com/"
 
@@ -64,6 +65,7 @@ class OpenAiProvider(Provider):
             total_tokens=total_tokens,
         )
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/exchange/tests/test_langfuse_wrapper.py

@@ -0,0 +1,46 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from exchange.langfuse_wrapper import observe_wrapper
+
+
+@pytest.fixture
+def mock_langfuse_context():
+    with patch("exchange.langfuse_wrapper.langfuse_context") as mock:
+        yield mock
+
+
+@patch("exchange.langfuse_wrapper.HAS_LANGFUSE_CREDENTIALS", True)
+def test_function_is_wrapped(mock_langfuse_context):
+    mock_observe = MagicMock(side_effect=lambda *args, **kwargs: lambda fn: fn)
+    mock_langfuse_context.observe = mock_observe
+
+    def original_function(x: int, y: int) -> int:
+        return x + y
+
+    # test function before we decorate it with
+    # @observe_wrapper("arg1", kwarg1="kwarg1")
+    assert not hasattr(original_function, "__wrapped__")
+
+    # ensure we args get passed along (e.g. @observe(capture_input=False, capture_output=False))
+    decorated_function = observe_wrapper("arg1", kwarg1="kwarg1")(original_function)
+    assert hasattr(decorated_function, "__wrapped__")
+    assert decorated_function.__wrapped__ is original_function, "Function is not properly wrapped"
+
+    assert decorated_function(2, 3) == 5
+    mock_observe.assert_called_once()
+    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):
+    mock_observe = MagicMock(return_value=lambda f: f)
+    mock_langfuse_context.observe = mock_observe
+
+    @observe_wrapper("arg1", kwarg1="kwarg1")
+    def hello() -> str:
+        return "Hello"
+
+    assert not hasattr(hello, "__wrapped__")
+    assert hello() == "Hello"
+
+    mock_observe.assert_not_called()