diff --git a/.gitignore b/.gitignore index f799b7221698..807d19890ab6 100644 --- a/.gitignore +++ b/.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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 1a5597e44382..deea382114a8 100644 --- a/CONTRIBUTING.md +++ b/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. diff --git a/README.md b/README.md index 1e7b0dc51a7e..f1d3cb81a333 100644 --- a/README.md +++ b/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]. diff --git a/docs/plugins/cli.md b/docs/plugins/cli.md index 5d27563c9191..7a3566d34474 100644 --- a/docs/plugins/cli.md +++ b/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:** diff --git a/justfile b/justfile index 27e813f9c8d0..96f51d6b8860 100644 --- a/justfile +++ b/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 diff --git a/packages/exchange/.env.langfuse.local b/packages/exchange/.env.langfuse.local new file mode 100644 index 000000000000..cdebcd7a16b2 --- /dev/null +++ b/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 diff --git a/packages/exchange/pyproject.toml b/packages/exchange/pyproject.toml index 8b30a9f3e1df..a25782f9f0db 100644 --- a/packages/exchange/pyproject.toml +++ b/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] diff --git a/packages/exchange/src/exchange/exchange.py b/packages/exchange/src/exchange/exchange.py index 05eda0df1f39..942bf78c669d 100644 --- a/packages/exchange/src/exchange/exchange.py +++ b/packages/exchange/src/exchange/exchange.py @@ -3,6 +3,7 @@ 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 @@ 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) diff --git a/packages/exchange/src/exchange/langfuse_wrapper.py b/packages/exchange/src/exchange/langfuse_wrapper.py new file mode 100644 index 000000000000..c8cec23eeec9 --- /dev/null +++ b/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 diff --git a/packages/exchange/src/exchange/providers/anthropic.py b/packages/exchange/src/exchange/providers/anthropic.py index a6a0c2262184..c98c6d432c99 100644 --- a/packages/exchange/src/exchange/providers/anthropic.py +++ b/packages/exchange/src/exchange/providers/anthropic.py @@ -7,6 +7,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 exchange.langfuse_wrapper import observe_wrapper ANTHROPIC_HOST = "https://api.anthropic.com/v1/messages" @@ -122,6 +123,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, diff --git a/packages/exchange/src/exchange/providers/bedrock.py b/packages/exchange/src/exchange/providers/bedrock.py index 1dd0ebadf4fd..cdc0c29c93df 100644 --- a/packages/exchange/src/exchange/providers/bedrock.py +++ b/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 exchange.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, diff --git a/packages/exchange/src/exchange/providers/databricks.py b/packages/exchange/src/exchange/providers/databricks.py index 052d78a67eec..b8f92dca9320 100644 --- a/packages/exchange/src/exchange/providers/databricks.py +++ b/packages/exchange/src/exchange/providers/databricks.py @@ -11,7 +11,7 @@ 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 @@ def get_usage(data: dict) -> Usage: total_tokens=total_tokens, ) + @observe_wrapper(as_type="generation") def complete( self, model: str, diff --git a/packages/exchange/src/exchange/providers/google.py b/packages/exchange/src/exchange/providers/google.py index e5f9312d1fc1..76ccd7a9cdf3 100644 --- a/packages/exchange/src/exchange/providers/google.py +++ b/packages/exchange/src/exchange/providers/google.py @@ -7,6 +7,8 @@ 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 @@ 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, diff --git a/packages/exchange/src/exchange/providers/groq.py b/packages/exchange/src/exchange/providers/groq.py index edd0945a165c..0f6472f887f3 100644 --- a/packages/exchange/src/exchange/providers/groq.py +++ b/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 @@ def get_usage(data: dict) -> Usage: total_tokens=total_tokens, ) + @observe_wrapper(as_type="generation") def complete( self, model: str, diff --git a/packages/exchange/src/exchange/providers/openai.py b/packages/exchange/src/exchange/providers/openai.py index 02bdcf4a3651..8701e5429c2c 100644 --- a/packages/exchange/src/exchange/providers/openai.py +++ b/packages/exchange/src/exchange/providers/openai.py @@ -14,6 +14,7 @@ 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 @@ def get_usage(data: dict) -> Usage: total_tokens=total_tokens, ) + @observe_wrapper(as_type="generation") def complete( self, model: str, diff --git a/packages/exchange/tests/test_langfuse_wrapper.py b/packages/exchange/tests/test_langfuse_wrapper.py new file mode 100644 index 000000000000..3218509785ad --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index c9c4501b5884..cf1a7609dcb6 100644 --- a/pyproject.toml +++ b/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" }] diff --git a/scripts/langfuse-docker-compose.yaml b/scripts/langfuse-docker-compose.yaml new file mode 100644 index 000000000000..1539321205cf --- /dev/null +++ b/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 diff --git a/scripts/setup_langfuse.sh b/scripts/setup_langfuse.sh new file mode 100755 index 000000000000..4807479598a3 --- /dev/null +++ b/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 diff --git a/src/goose/cli/main.py b/src/goose/cli/main.py index aa345b60a014..fedae331eb2b 100644 --- a/src/goose/cli/main.py +++ b/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]: diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py index 666901b6a59f..8cdb2f646ecf 100644 --- a/src/goose/cli/session.py +++ b/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 @@ def __init__( 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 @@ def __init__( 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 @@ def run(self, new_session: bool = True) -> None: 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") diff --git a/tests/cli/test_main.py b/tests/cli/test_main.py index 9f6f2f66ffad..38b38920f537 100644 --- a/tests/cli/test_main.py +++ b/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()