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

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

.gitignore

@@ -103,6 +103,9 @@ celerybeat.pid
 .env.*
 .venv
 
+# exception for local langfuse init vars
+!**/packages/exchange/.env.langfuse.local 
+
 # Spyder project settings
 .spyderproject
 .spyproject

CONTRIBUTING.md

@@ -48,6 +48,21 @@ or, as a shortcut,
 just test
 ```
 
+### Enable traces in Goose with [locally hosted Langfuse](https://langfuse.com/docs/deployment/self-host)
+> [!NOTE]
+> This integration is experimental and we don't currently have integration tests for it.
+ 
+Developers can use locally hosted Langfuse tracing by applying the custom `observe_wrapper` decorator defined in `packages/exchange/src/langfuse_wrapper.py` to functions for automatic integration with Langfuse. 
+
+- Run `just langfuse-server` to start your local Langfuse server. It requires Docker.
+- Go to http://localhost:3000 and log in with the default email/password output by the shell script (values can also be found in the `.env.langfuse.local` file).
+- Run Goose with the --tracing flag enabled i.e., `goose session start --tracing`
+- View your traces at http://localhost:3000
+
+To extend tracing to additional functions, import `from exchange.langfuse_wrapper import observe_wrapper` and use the `observe_wrapper()` decorator on functions you wish to enable tracing for. `observe_wrapper` functions the same way as Langfuse's observe decorator. 
+
+Read more about Langfuse's decorator-based tracing [here](https://langfuse.com/docs/sdk/python/decorators).
+
 ## Exchange
 
 The lower level generation behind goose is powered by the [`exchange`][ai-exchange] package, also in this repo.

README.md

@@ -150,6 +150,34 @@ goose session resume
 
 To see more documentation on the CLI commands currently available to Goose check out the documentation [here][cli]. If you’d like to develop your own CLI commands for Goose, check out the [Contributing document][contributing].
 
+### Tracing with Langfuse
+> [!NOTE]
+> This Langfuse integration is experimental and we don't currently have integration tests for it.
+
+The exchange package provides a [Langfuse](https://langfuse.com/) wrapper module. The wrapper serves to initialize Langfuse appropriately if the Langfuse server is running locally and otherwise to skip applying the Langfuse observe descorators.
+
+#### Start your local Langfuse server
+
+Run `just langfuse-server` to start your local Langfuse server. It requires Docker.
+
+Read more about local Langfuse deployments [here](https://langfuse.com/docs/deployment/local).
+
+#### Exchange and Goose integration
+
+Import `from exchange.langfuse_wrapper import observe_wrapper` and use the `observe_wrapper()` decorator on functions you wish to enable tracing for. `observe_wrapper` functions the same way as Langfuse's observe decorator.
+
+Read more about Langfuse's decorator-based tracing [here](https://langfuse.com/docs/sdk/python/decorators).
+
+In Goose, initialization requires certain environment variables to be present:
+
+- `LANGFUSE_PUBLIC_KEY`: Your Langfuse public key
+- `LANGFUSE_SECRET_KEY`: Your Langfuse secret key
+- `LANGFUSE_BASE_URL`: The base URL of your Langfuse instance
+
+By default your local deployment and Goose will use the values in `.env.langfuse.local`.
+
+
+
 ### Next steps
 
 Learn how to modify your Goose profiles.yaml file to add and remove functionality (toolkits) and providing context to get the most out of Goose in our [Getting Started Guide][getting-started].

docs/plugins/cli.md

@@ -19,11 +19,13 @@ Lists the version of Goose and any associated plugins.
 
 **Usage:**
 ```sh
-  goose session start [--profile PROFILE] [--plan PLAN]
+  goose session start [--profile PROFILE] [--plan PLAN] [--log-level [DEBUG|INFO|WARNING|ERROR|CRITICAL]] [--tracing]
 ```
 
 Starts a new Goose session.
 
+If you want to enable locally hosted Langfuse tracing, pass the --tracing flag after starting your local Langfuse server as outlined in the [Contributing Guide's][contributing] Development guidelines.
+
 #### `resume`
 
 **Usage:**

justfile

@@ -74,3 +74,6 @@ tag-push:
 # get commit messages for a release
 release-notes:
   git log --pretty=format:"- %s" v$(just tag_version)..HEAD
+
+langfuse-server:
+  ./scripts/setup_langfuse.sh

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()

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "click>=8.1.7",
     "prompt-toolkit>=3.0.47",
     "keyring>=25.4.1",
+    "langfuse>=2.38.2",
 ]
 author = [{ name = "Block", email = "ai-oss-tools@block.xyz" }]
 packages = [{ include = "goose", from = "src" }]

scripts/langfuse-docker-compose.yaml

@@ -0,0 +1,46 @@
+services:
+  langfuse-server:
+    image: langfuse/langfuse:2
+    depends_on:
+      db:
+        condition: service_healthy
+    ports:
+      - "3000:3000"
+    environment:
+      - DATABASE_URL=postgresql://postgres:postgres@db:5432/postgres
+      - NEXTAUTH_SECRET=mysecret
+      - SALT=mysalt
+      - ENCRYPTION_KEY=0000000000000000000000000000000000000000000000000000000000000000 # generate via `openssl rand -hex 32`
+      - NEXTAUTH_URL=http://localhost:3000
+      - TELEMETRY_ENABLED=${TELEMETRY_ENABLED:-true}
+      - LANGFUSE_ENABLE_EXPERIMENTAL_FEATURES=${LANGFUSE_ENABLE_EXPERIMENTAL_FEATURES:-false}
+      - LANGFUSE_INIT_ORG_ID=${LANGFUSE_INIT_ORG_ID:-}
+      - LANGFUSE_INIT_ORG_NAME=${LANGFUSE_INIT_ORG_NAME:-}
+      - LANGFUSE_INIT_PROJECT_ID=${LANGFUSE_INIT_PROJECT_ID:-}
+      - LANGFUSE_INIT_PROJECT_NAME=${LANGFUSE_INIT_PROJECT_NAME:-}
+      - LANGFUSE_INIT_PROJECT_PUBLIC_KEY=${LANGFUSE_INIT_PROJECT_PUBLIC_KEY:-}
+      - LANGFUSE_INIT_PROJECT_SECRET_KEY=${LANGFUSE_INIT_PROJECT_SECRET_KEY:-}
+      - LANGFUSE_INIT_USER_EMAIL=${LANGFUSE_INIT_USER_EMAIL:-}
+      - LANGFUSE_INIT_USER_NAME=${LANGFUSE_INIT_USER_NAME:-}
+      - LANGFUSE_INIT_USER_PASSWORD=${LANGFUSE_INIT_USER_PASSWORD:-}
+
+  db:
+    image: postgres
+    restart: always
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 3s
+      timeout: 3s
+      retries: 10
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+      - POSTGRES_DB=postgres
+    ports:
+      - 5432:5432
+    volumes:
+      - database_data:/var/lib/postgresql/data
+
+volumes:
+  database_data:
+    driver: local

scripts/setup_langfuse.sh

@@ -0,0 +1,99 @@
+#!/bin/bash
+
+# setup_langfuse.sh
+#
+# This script sets up and runs Langfuse locally for development and testing purposes.
+#
+# Key functionalities:
+# 1. Downloads the latest docker-compose.yaml from the Langfuse repository
+# 2. Starts Langfuse using Docker Compose with default initialization variables
+# 3. Waits for the service to be available
+# 4. Launches a browser to open the local Langfuse UI
+# 5. Prints login credentials from the environment file
+#
+# Usage:
+#   ./setup_langfuse.sh
+#
+# Requirements:
+#   - Docker
+#   - curl
+#   - A .env.langfuse.local file in the env directory
+#
+# Note: This script is intended for local development use only.
+
+set -e
+
+SCRIPT_DIR=$(realpath "$(dirname "${BASH_SOURCE[0]}")")
+LANGFUSE_DOCKER_COMPOSE_URL="https://raw.githubusercontent.com/langfuse/langfuse/main/docker-compose.yml"
+LANGFUSE_DOCKER_COMPOSE_FILE="langfuse-docker-compose.yaml"
+LANGFUSE_ENV_FILE="$SCRIPT_DIR/../packages/exchange/.env.langfuse.local"
+
+check_dependencies() {
+    local dependencies=("curl" "docker")
+    local missing_dependencies=()
+
+    for cmd in "${dependencies[@]}"; do
+        if ! command -v "$cmd" &> /dev/null; then
+            missing_dependencies+=("$cmd")
+        fi
+    done
+
+    if [ ${#missing_dependencies[@]} -ne 0 ]; then
+        echo "Missing dependencies: ${missing_dependencies[*]}"
+        exit 1
+    fi
+}
+
+download_docker_compose() {
+    if ! curl --fail --location --output "$SCRIPT_DIR/langfuse-docker-compose.yaml" "$LANGFUSE_DOCKER_COMPOSE_URL"; then
+        echo "Failed to download docker-compose file from $LANGFUSE_DOCKER_COMPOSE_URL"
+        exit 1
+    fi
+}
+
+start_docker_compose() {
+    docker compose --env-file "$LANGFUSE_ENV_FILE" -f "$LANGFUSE_DOCKER_COMPOSE_FILE" up --detach
+}
+
+wait_for_service() {
+    echo "Waiting for Langfuse to start..."
+    local retries=10
+    local count=0
+    until curl --silent http://localhost:3000 > /dev/null; do
+        ((count++))
+        if [ "$count" -ge "$retries" ]; then
+            echo "Max retries reached. Langfuse did not start in time."
+            exit 1
+        fi
+        sleep 1
+    done
+    echo "Langfuse is now available!"
+}
+
+launch_browser() {
+    if [[ "$OSTYPE" == "linux-gnu"* ]]; then
+        xdg-open "http://localhost:3000"
+    elif [[ "$OSTYPE" == "darwin"* ]]; then
+        open "http://localhost:3000"
+    else
+        echo "Please open http://localhost:3000 to view Langfuse traces."
+    fi
+}
+
+print_login_variables() {
+    if [ -f "$LANGFUSE_ENV_FILE" ]; then
+        echo "If not already logged in use the following credentials to log in:"
+        grep -E "LANGFUSE_INIT_USER_EMAIL|LANGFUSE_INIT_USER_PASSWORD" "$LANGFUSE_ENV_FILE"
+    else
+        echo "Langfuse environment file with local credentials not found."
+    fi
+}
+
+check_dependencies
+pushd "$SCRIPT_DIR" > /dev/null
+download_docker_compose
+start_docker_compose
+wait_for_service
+print_login_variables
+launch_browser
+popd > /dev/null

src/goose/cli/main.py

@@ -138,8 +138,11 @@ def get_session_files() -> dict[str, Path]:
 @click.argument("name", required=False, shell_complete=autocomplete_session_files)
 @click.option("--profile")
 @click.option("--plan", type=click.Path(exists=True))
-@click.option("--log-level", type=LOG_CHOICE, default="INFO")
-def session_start(name: Optional[str], profile: str, log_level: str, plan: Optional[str] = None) -> None:
+@click.option("--log-level", type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]), default="INFO")
+@click.option("--tracing", is_flag=True, required=False)
+def session_start(
+    name: Optional[str], profile: str, log_level: str, plan: Optional[str] = None, tracing: bool = False
+) -> None:
     """Start a new goose session"""
     if plan:
         yaml = YAML()
@@ -147,8 +150,12 @@ def session_start(name: Optional[str], profile: str, log_level: str, plan: Optio
             _plan = yaml.load(f)
     else:
         _plan = None
-    session = Session(name=name, profile=profile, plan=_plan, log_level=log_level)
-    session.run()
+
+    try:
+        session = Session(name=name, profile=profile, plan=_plan, log_level=log_level, tracing=tracing)
+        session.run()
+    except RuntimeError as e:
+        print(f"[red]Error: {e}")
 
 
 def parse_args(ctx: click.Context, param: click.Parameter, value: str) -> dict[str, str]:

src/goose/cli/session.py

@@ -1,8 +1,11 @@
+import logging
 import traceback
 from pathlib import Path
 from typing import Optional
 
+from langfuse.decorators import langfuse_context
 from exchange import Message, Text, ToolResult, ToolUse
+from exchange.langfuse_wrapper import observe_wrapper, auth_check
 from rich import print
 from rich.markdown import Markdown
 from rich.panel import Panel
@@ -62,6 +65,7 @@ class Session:
         profile: Optional[str] = None,
         plan: Optional[dict] = None,
         log_level: Optional[str] = "INFO",
+        tracing: bool = False,
         **kwargs: dict[str, any],
     ) -> None:
         if name is None:
@@ -72,6 +76,18 @@ class Session:
         self.prompt_session = GoosePromptSession()
         self.status_indicator = Status("", spinner="dots")
         self.notifier = SessionNotifier(self.status_indicator)
+        if not tracing:
+            logging.getLogger("langfuse").setLevel(logging.ERROR)
+        else:
+            langfuse_auth = auth_check()
+            if langfuse_auth:
+                print("Local Langfuse initialized. View your traces at http://localhost:3000")
+            else:
+                raise RuntimeError(
+                    "You passed --tracing, but a Langfuse object was not found in the current context. "
+                    "Please initialize the local Langfuse server and restart Goose."
+                )
+        langfuse_context.configure(enabled=tracing)
 
         self.exchange = create_exchange(profile=load_profile(profile), notifier=self.notifier)
         setup_logging(log_file_directory=LOG_PATH, log_level=log_level)
@@ -189,6 +205,7 @@ class Session:
         self._remove_empty_session()
         self._log_cost()
 
+    @observe_wrapper()
     def reply(self) -> None:
         """Reply to the last user message, calling tools as needed"""
         self.status_indicator.update("responding")

tests/cli/test_main.py

@@ -34,7 +34,9 @@ def test_session_start_command_with_session_name(mock_session):
     mock_session_class, mock_session_instance = mock_session
     runner = CliRunner()
     runner.invoke(goose_cli, ["session", "start", "session1", "--profile", "default"])
-    mock_session_class.assert_called_once_with(name="session1", profile="default", plan=None, log_level="INFO")
+    mock_session_class.assert_called_once_with(
+        name="session1", profile="default", plan=None, log_level="INFO", tracing=False
+    )
     mock_session_instance.run.assert_called_once()