From 1ae33529011f23f6bc9653b2167bb6a65aa4afac Mon Sep 17 00:00:00 2001 From: Adrian Cole Date: Wed, 18 Sep 2024 18:10:32 +0800 Subject: [PATCH 1/5] feat: add basic opentelemetry tracing support --- .otel.env | 5 ++ CONTRIBUTING.md | 21 ++++++- justfile | 3 + pyproject.toml | 6 ++ src/goose/cli/session.py | 52 +++++++++++------ tests/cli/test_session_otel.py | 101 +++++++++++++++++++++++++++++++++ 6 files changed, 170 insertions(+), 18 deletions(-) create mode 100644 .otel.env create mode 100644 tests/cli/test_session_otel.py diff --git a/.otel.env b/.otel.env new file mode 100644 index 00000000..8470429e --- /dev/null +++ b/.otel.env @@ -0,0 +1,5 @@ +OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 +OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf +OTEL_SERVICE_NAME=goose +OTEL_METRIC_EXPORT_INTERVAL=100 +OTEL_BSP_SCHEDULE_DELAY=100 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7a4edc4a..a2770aab 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -108,6 +108,24 @@ Additions to the [developer toolkit][developer] change the core performance, and This project follows the [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) specification for PR titles. Conventional Commits make it easier to understand the history of a project and facilitate automation around versioning and changelog generation. +## Running with OpenTelemetry + +Goose uses HTTP to access model providers. If you are receiving failures, it can be helpful to see traces +of the underlying HTTP requests. For example, a 404 could be indicative of an incorrect URL or a missing model. + +First, ensure you have an OpenTelemetry compatible collector listening on port 4318, such as [otel-tui][otel-tui]. + +```bash +brew install ymtdzzz/tap/otel-tui +otel-tui +``` + +Then, start goose like this: +```bash +uv run dotenv -f ./.otel.env run -- opentelemetry-instrument goose session start +# or `just otel-goose session start` +``` + [issues]: https://github.com/square/goose/issues [goose-plugins]: https://github.com/square/goose-plugins [ai-exchange]: https://github.com/square/exchange @@ -115,4 +133,5 @@ This project follows the [Conventional Commits](https://www.conventionalcommits. [uv]: https://docs.astral.sh/uv/ [ruff]: https://docs.astral.sh/ruff/ [just]: https://github.com/casey/just -[toolkits]: docs/docs/toolkits.md \ No newline at end of file +[toolkits]: docs/docs/toolkits.md +[otel-tui]: https://github.com/ymtdzzz/otel-tui diff --git a/justfile b/justfile index 4d3ac634..ce69ea6c 100644 --- a/justfile +++ b/justfile @@ -20,3 +20,6 @@ coverage *FLAGS: docs: cd docs && uv sync && uv run mkdocs serve + +otel-goose *FLAGS: + uv run dotenv -f ./.otel.env run -- opentelemetry-instrument goose {{FLAGS}} diff --git a/pyproject.toml b/pyproject.toml index b51e5742..c6a61407 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "ai-exchange>=0.9.0", "click>=8.1.7", "prompt-toolkit>=3.0.47", + "opentelemetry-api>=1.27.0", ] author = [{ name = "Block", email = "ai-oss-tools@block.xyz" }] packages = [{ include = "goose", from = "src" }] @@ -61,6 +62,11 @@ dev-dependencies = [ "mkdocs-redirects>=1.2.1", "mkdocs-include-markdown-plugin>=6.2.2", "mkdocs-callouts>=1.14.0", + "opentelemetry-sdk>=1.27.0", + "opentelemetry-exporter-otlp-proto-http>=1.27.0", + "opentelemetry-distro>=0.48b0", + "opentelemetry-instrumentation-httpx>=0.48b0", + "python-dotenv[cli]>=1.0.1" ] diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py index 1f72a302..28635540 100644 --- a/src/goose/cli/session.py +++ b/src/goose/cli/session.py @@ -3,6 +3,8 @@ from typing import Any, Dict, List, Optional from exchange import Message, ToolResult, ToolUse, Text +from opentelemetry import trace +from opentelemetry.trace.status import Status as OtelStatus, StatusCode as OtelStatusCode from prompt_toolkit.shortcuts import confirm from rich import print from rich.console import RenderableType @@ -90,11 +92,13 @@ def __init__( name: Optional[str] = None, profile: Optional[str] = None, plan: Optional[dict] = None, + tracer: trace.Tracer = trace.get_tracer(__name__), **kwargs: Dict[str, Any], ) -> None: self.name = name self.status_indicator = Status("", spinner="dots") self.notifier = SessionNotifier(self.status_indicator) + self.tracer = tracer self.exchange = build_exchange(profile=load_profile(profile), notifier=self.notifier) @@ -150,23 +154,37 @@ def run(self) -> None: """ message = self.process_first_message() while message: # Loop until no input (empty string). - self.notifier.start() - try: - self.exchange.add(message) - self.reply() # Process the user message. - except KeyboardInterrupt: - self.interrupt_reply() - except Exception: - # rewind to right before the last user message - self.exchange.rewind() - print(traceback.format_exc()) - print( - "\n[red]The error above was an exception we were not able to handle.\n\n[/]" - + "These errors are often related to connection or authentication\n" - + "We've removed the conversation up to the most recent user message" - + " - [yellow]depending on the error you may be able to continue[/]" - ) - self.notifier.stop() + # Start a span with the default tracer. This is present with opentelemetry-instrument + with self.tracer.start_as_current_span(message.role) as span: + span.set_attribute("goose.role", message.role) + span.set_attribute("goose.id", message.id) + # For starters, only add to the trace the first text content + first_content = message.content[0] if message.content else None + if isinstance(first_content, Text): + span.set_attribute("goose.text", first_content.text) + + self.notifier.start() + try: + self.exchange.add(message) + self.reply() # Process the user message. + except KeyboardInterrupt: + span.set_status(OtelStatus(OtelStatusCode.ERROR, "KeyboardInterrupt")) + self.interrupt_reply() + except Exception as e: + # rewind to right before the last user message + self.exchange.rewind() + span.set_status(OtelStatus(OtelStatusCode.ERROR, str(e))) + span.record_exception(e) + print(traceback.format_exc()) + print( + "\n[red]The error above was an exception we were not able to handle.\n\n[/]" + + "These errors are often related to connection or authentication\n" + + "We've removed the conversation up to the most recent user message" + + " - [yellow]depending on the error you may be able to continue[/]" + ) + finally: + self.notifier.stop() + span.end() print() # Print a newline for separation. user_input = self.prompt_session.get_user_input() diff --git a/tests/cli/test_session_otel.py b/tests/cli/test_session_otel.py new file mode 100644 index 00000000..dec8a3a2 --- /dev/null +++ b/tests/cli/test_session_otel.py @@ -0,0 +1,101 @@ +import pytest +from httpx import HTTPStatusError +from unittest.mock import MagicMock, patch +from exchange import Message, Text +from goose.cli.session import Session +from goose.cli.prompt.goose_prompt_session import GoosePromptSession +from goose.cli.prompt.user_input import PromptAction, UserInput +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +@pytest.fixture +def memory_exporter(): + yield InMemorySpanExporter() + + +@pytest.fixture +def create_session_with_mock_configs(mock_sessions_path, exchange_factory, profile_factory, memory_exporter): + # Create a tracer with the same memory exporter as test assertions use. We do this for each + # test to ensure they can run in parallel without interfering with each other. + tracer_provider = TracerProvider() + memory_exporter = memory_exporter + span_processor = SimpleSpanProcessor(memory_exporter) + tracer_provider.add_span_processor(span_processor) + + with patch("goose.cli.session.build_exchange", return_value=exchange_factory()), patch( + "goose.cli.session.load_profile", return_value=profile_factory() + ), patch("goose.cli.session.SessionNotifier") as mock_session_notifier, patch( + "goose.cli.session.load_provider", return_value="provider" + ): + mock_session_notifier.return_value = MagicMock() + + def create_session(session_attributes: dict = {}): + return Session(**session_attributes, tracer=tracer_provider.get_tracer(__name__)) + + yield create_session + + +def test_trace_run(create_session_with_mock_configs, memory_exporter): + session = create_session_with_mock_configs() + + message = Message(role="user", id="abracadabra", content=[Text("List the files in this directory")]) + + # Call the run function, for a single message which results in an exit. + with patch.object(Session, "process_first_message", return_value=message), patch.object( + Session, "reply", return_value=None + ), patch.object( + GoosePromptSession, "get_user_input", return_value=UserInput(action=PromptAction.EXIT) + ), patch.object(Session, "save_session", return_value=None): + session.run() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == message.role + assert span.attributes == {"goose.id": message.id, "goose.role": message.role, "goose.text": message.text} + + +def test_trace_run_interrupted(create_session_with_mock_configs, memory_exporter): + session = create_session_with_mock_configs() + + message = Message(role="user", id="abracadabra", content=[Text("List the files in this directory")]) + + # Call the run function, for a single message which results in an exit. + with patch.object(Session, "process_first_message", return_value=message), patch.object( + Session, "reply", side_effect=KeyboardInterrupt + ), patch.object( + GoosePromptSession, "get_user_input", return_value=UserInput(action=PromptAction.EXIT) + ), patch.object(Session, "save_session", return_value=None): + session.run() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == message.role + assert span.attributes == {"goose.id": message.id, "goose.role": message.role, "goose.text": message.text} + assert span.status.is_ok is False + assert span.status.description == "KeyboardInterrupt" + + +def test_trace_run_error(create_session_with_mock_configs, memory_exporter): + session = create_session_with_mock_configs() + + message = Message(role="user", id="abracadabra", content=[Text("List the files in this directory")]) + + # Call the run function, for a single message which results in an HTTP error. + with patch.object(Session, "process_first_message", return_value=message), patch.object( + Session, "reply", side_effect=HTTPStatusError("HTTP error", request=None, response=None) + ), patch.object( + GoosePromptSession, "get_user_input", return_value=UserInput(action=PromptAction.EXIT) + ), patch.object(Session, "save_session", return_value=None): + session.run() + + spans = memory_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == message.role + assert span.attributes == {"goose.id": message.id, "goose.role": message.role, "goose.text": message.text} + assert span.status.is_ok is False + assert span.status.description == "HTTP error" From d7596b0d506478874c994114264da3f27adb2a29 Mon Sep 17 00:00:00 2001 From: Adrian Cole <64215+codefromthecrypt@users.noreply.github.com> Date: Wed, 18 Sep 2024 18:21:20 +0800 Subject: [PATCH 2/5] Update tests/cli/test_session_otel.py --- tests/cli/test_session_otel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cli/test_session_otel.py b/tests/cli/test_session_otel.py index dec8a3a2..02e7e5ac 100644 --- a/tests/cli/test_session_otel.py +++ b/tests/cli/test_session_otel.py @@ -62,7 +62,7 @@ def test_trace_run_interrupted(create_session_with_mock_configs, memory_exporter message = Message(role="user", id="abracadabra", content=[Text("List the files in this directory")]) - # Call the run function, for a single message which results in an exit. + # Call the run function, for a single message which results in an interrupt. with patch.object(Session, "process_first_message", return_value=message), patch.object( Session, "reply", side_effect=KeyboardInterrupt ), patch.object( From 510175427365b7d9375ecf964d0a08abba8af9ff Mon Sep 17 00:00:00 2001 From: Adrian Cole Date: Thu, 19 Sep 2024 15:53:59 +0800 Subject: [PATCH 3/5] feedback Signed-off-by: Adrian Cole --- .otel.env | 5 ----- CONTRIBUTING.md | 5 ++--- justfile | 2 +- src/goose/cli/session.py | 25 ++++++++++++++++--------- 4 files changed, 19 insertions(+), 18 deletions(-) delete mode 100644 .otel.env diff --git a/.otel.env b/.otel.env deleted file mode 100644 index 8470429e..00000000 --- a/.otel.env +++ /dev/null @@ -1,5 +0,0 @@ -OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 -OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf -OTEL_SERVICE_NAME=goose -OTEL_METRIC_EXPORT_INTERVAL=100 -OTEL_BSP_SCHEDULE_DELAY=100 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a2770aab..b4d1f252 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -116,13 +116,12 @@ of the underlying HTTP requests. For example, a 404 could be indicative of an in First, ensure you have an OpenTelemetry compatible collector listening on port 4318, such as [otel-tui][otel-tui]. ```bash -brew install ymtdzzz/tap/otel-tui -otel-tui +docker run --rm -it --name otel-tui -p 4318:4318 ymtdzzz/otel-tui ``` Then, start goose like this: ```bash -uv run dotenv -f ./.otel.env run -- opentelemetry-instrument goose session start +uv run opentelemetry-instrument --service_name goose --exporter_otlp_protocol http/protobuf goose session start # or `just otel-goose session start` ``` diff --git a/justfile b/justfile index ce69ea6c..889c39f8 100644 --- a/justfile +++ b/justfile @@ -22,4 +22,4 @@ docs: cd docs && uv sync && uv run mkdocs serve otel-goose *FLAGS: - uv run dotenv -f ./.otel.env run -- opentelemetry-instrument goose {{FLAGS}} + uv run opentelemetry-instrument --service_name goose --exporter_otlp_protocol http/protobuf goose {{FLAGS}} diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py index 28635540..c1875efd 100644 --- a/src/goose/cli/session.py +++ b/src/goose/cli/session.py @@ -92,12 +92,15 @@ def __init__( name: Optional[str] = None, profile: Optional[str] = None, plan: Optional[dict] = None, - tracer: trace.Tracer = trace.get_tracer(__name__), + tracer: trace.Tracer = trace.get_tracer("goose"), **kwargs: Dict[str, Any], ) -> None: self.name = name self.status_indicator = Status("", spinner="dots") self.notifier = SessionNotifier(self.status_indicator) + + # Set the tracer as a field in session, as opposed to a module variable + # so that tests can swap this out and safely run in parallel. self.tracer = tracer self.exchange = build_exchange(profile=load_profile(profile), notifier=self.notifier) @@ -154,20 +157,24 @@ def run(self) -> None: """ message = self.process_first_message() while message: # Loop until no input (empty string). - # Start a span with the default tracer. This is present with opentelemetry-instrument - with self.tracer.start_as_current_span(message.role) as span: - span.set_attribute("goose.role", message.role) - span.set_attribute("goose.id", message.id) - # For starters, only add to the trace the first text content - first_content = message.content[0] if message.content else None - if isinstance(first_content, Text): - span.set_attribute("goose.text", first_content.text) + span_attributes = { + "goose.role": message.role, + "goose.id": message.id, + } + + # For starters, only add to the trace the first text content + first_content = message.content[0] if message.content else None + if isinstance(first_content, Text): + span_attributes["goose.text"] = first_content.text + with self.tracer.start_as_current_span(message.role, attributes=span_attributes) as span: self.notifier.start() try: self.exchange.add(message) self.reply() # Process the user message. except KeyboardInterrupt: + # TODO: should we make interrupting an error? If not, how should we mark this? + # span as it was interrupted? span.set_status(OtelStatus(OtelStatusCode.ERROR, "KeyboardInterrupt")) self.interrupt_reply() except Exception as e: From 081b9df9a2fc6fd5266f79054506ae81429b5f24 Mon Sep 17 00:00:00 2001 From: Adrian Cole Date: Thu, 19 Sep 2024 16:06:51 +0800 Subject: [PATCH 4/5] double-end Signed-off-by: Adrian Cole --- src/goose/cli/session.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/goose/cli/session.py b/src/goose/cli/session.py index c1875efd..ac6fabf1 100644 --- a/src/goose/cli/session.py +++ b/src/goose/cli/session.py @@ -189,9 +189,7 @@ def run(self) -> None: + "We've removed the conversation up to the most recent user message" + " - [yellow]depending on the error you may be able to continue[/]" ) - finally: - self.notifier.stop() - span.end() + self.notifier.stop() print() # Print a newline for separation. user_input = self.prompt_session.get_user_input() From 0796b0776df2a53b341db292c96a395d5ea5db72 Mon Sep 17 00:00:00 2001 From: Adrian Cole <64215+codefromthecrypt@users.noreply.github.com> Date: Thu, 19 Sep 2024 16:08:35 +0800 Subject: [PATCH 5/5] stale dep --- pyproject.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c6a61407..1d6953cb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,6 @@ dev-dependencies = [ "opentelemetry-exporter-otlp-proto-http>=1.27.0", "opentelemetry-distro>=0.48b0", "opentelemetry-instrumentation-httpx>=0.48b0", - "python-dotenv[cli]>=1.0.1" ]