feat: add local langfuse tracing option (#106)

.github/workflows/ci.yaml

@@ -28,7 +28,7 @@ jobs:
           uv run pytest tests -m 'not integration'
 
   goose:
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-latest 
 
     steps:
       - uses: actions/checkout@v4
@@ -48,6 +48,27 @@ jobs:
         run: |
           uv run pytest tests -m 'not integration'
 
+  langfuse-wrapper:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install UV
+        run: curl -LsSf https://astral.sh/uv/install.sh | sh
+
+      - name: Source Cargo Environment
+        run: source $HOME/.cargo/env
+
+      - name: Ruff
+        run: |
+          uvx ruff check packages/langfuse-wrapper
+          uvx ruff format packages/langfuse-wrapper --check
+
+      - name: Run tests
+        working-directory: ./packages/langfuse-wrapper
+        run: |
+          uv run pytest tests -m 'not integration'
 
   # This runs integration tests of the OpenAI API, using Ollama to host models.
   # This lets us test PRs from forks which can't access secrets like API keys.

.gitignore

@@ -125,3 +125,6 @@ docs/docs/reference
 
 # uv lock file
 uv.lock
+
+# langfuse docker file
+**/packages/langfuse-wrapper/scripts/docker-compose.yaml

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-wrapper"   
 ]
 
 [tool.hatch.build.targets.wheel]
@@ -47,3 +49,10 @@ ai-exchange = "exchange:module_name"
 markers = [
     "integration: marks tests that need to authenticate (deselect with '-m \"not integration\"')",
 ]
+
+[tool.uv.sources]
+langfuse-wrapper = { workspace = true}
+
+[tool.uv.workspace]
+members = ["../langfuse-wrapper"]
+

packages/exchange/src/exchange/exchange.py

@@ -14,6 +14,7 @@
 from exchange.providers import Provider, Usage
 from exchange.tool import Tool
 from exchange.token_usage_collector import _token_usage_collector
+from langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 
 def validate_tool_output(output: str) -> None:
@@ -127,6 +128,7 @@ def reply(self, max_tool_use: int = 128) -> Message:
 
         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/providers/anthropic.py

@@ -8,6 +8,7 @@
 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 langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 ANTHROPIC_HOST = "https://api.anthropic.com/v1/messages"
 
@@ -123,6 +124,7 @@ def messages_to_anthropic_spec(messages: List[Message]) -> List[Dict[str, Any]]:
             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 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 langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 SERVICE = "bedrock-runtime"
 UTC = timezone.utc
@@ -175,6 +176,7 @@ def from_env(cls: Type["BedrockProvider"]) -> "BedrockProvider":
         )
         return cls(client=client)
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

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

@@ -13,7 +13,7 @@
     tools_to_openai_spec,
 )
 from exchange.tool import Tool
-
+from langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 retry_procedure = retry(
     wait=wait_fixed(2),
@@ -69,6 +69,7 @@ def get_usage(data: dict) -> Usage:
             total_tokens=total_tokens,
         )
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

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

@@ -8,6 +8,7 @@
 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
+from langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 GOOGLE_HOST = "https://generativelanguage.googleapis.com/v1beta"
 
@@ -121,6 +122,7 @@ def messages_to_google_spec(messages: List[Message]) -> List[Dict[str, Any]]:
 
         return messages_spec
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

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

@@ -15,6 +15,7 @@
 from exchange.tool import Tool
 from tenacity import retry, wait_fixed, stop_after_attempt
 from exchange.providers.utils import retry_if_status
+from langfuse_wrapper.langfuse_wrapper import observe_wrapper
 
 OPENAI_HOST = "https://api.openai.com/"
 
@@ -65,6 +66,7 @@ def get_usage(data: dict) -> Usage:
             total_tokens=total_tokens,
         )
 
+    @observe_wrapper(as_type="generation")
     def complete(
         self,
         model: str,

packages/langfuse-wrapper/.ruff.toml

@@ -0,0 +1,2 @@
+lint.select = ["E", "W", "F", "N"]
+line-length = 120

packages/langfuse-wrapper/README.md

@@ -0,0 +1,28 @@
+# Langfuse Wrapper
+
+This package provides a wrapper for [Langfuse](https://langfuse.com/). The wrapper serves to initialize Langfuse appropriately if the Langfuse server is running locally and otherwise to skip applying the Langfuse observe descorators.
+
+**Note: This Langfuse integration is experimental and we don't currently have integration tests for it.**
+
+
+## Usage
+
+### Start your local Langfuse server
+
+Run `setup_langfuse.sh` 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 langfuse_wrapper.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/.env.langfuse.local`.

packages/langfuse-wrapper/env/.env.langfuse.local

@@ -0,0 +1,16 @@
+# These variables are default initialization variables for the Langfuse server
+LANGFUSE_INIT_PROJECT_NAME=goose-local
+LANGFUSE_INIT_PROJECT_PUBLIC_KEY=publickey-local
+LANGFUSE_INIT_PROJECT_SECRET_KEY=secretkey-local
+[email protected]
+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/langfuse-wrapper/pyproject.toml

@@ -0,0 +1,28 @@
+[project]
+name = "langfuse-wrapper"
+version = "0.1.0"
+description = "A wrapper for Langfuse integration"
+readme = "README.md"
+requires-python = ">=3.10"
+author = [{ name = "Block", email = "[email protected]" }]
+packages = [{ include = "langfuse_wrapper", from = "src" }]
+
+dependencies = [
+    "langfuse>=2.38.2",
+    "python-dotenv>=1.0.1"
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/langfuse_wrapper"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.uv]
+dev-dependencies = ["pytest>=8.3.2"]
+
+[tool.pytest.ini_options]
+markers = [
+    "integration: marks tests that need to authenticate (deselect with '-m \"not integration\"')",
+]

packages/langfuse-wrapper/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="docker-compose.yaml"
+LANGFUSE_ENV_FILE="$SCRIPT_DIR/../env/.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/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

packages/langfuse-wrapper/src/langfuse_wrapper/__init__.py

@@ -0,0 +1,3 @@
+from langfuse_wrapper.langfuse_wrapper import observe_wrapper  # noqa
+
+module_name = "langfuse-wrapper"