Add observability

griptape-ai · Jul 16, 2024 · f223a5d · f223a5d
1 parent 71a499d
commit f223a5d
Show file tree

Hide file tree

Showing 41 changed files with 1,777 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Native function calling support to `OpenAiChatPromptDriver`, `AzureOpenAiChatPromptDriver`, `AnthropicPromptDriver`, `AmazonBedrockPromptDriver`, `GooglePromptDriver`, and `CoherePromptDriver`.
 - `OllamaEmbeddingDriver` for generating embeddings with Ollama.
 - `GriptapeCloudKnowledgeBaseVectorStoreDriver` to query Griptape Cloud Knowledge Bases.
+- `BaseObservabilityDriver` as the base class for all Observability Drivers.
+- `DummyObservabilityDriver` as a no-op Observability Driver.
+- `OpenTelemetryObservabilityDriver` for sending observability data to an open telemetry collector or vendor.
+- `GriptapeCloudObservabilityDriver` for sending observability data to Griptape Cloud.
+- `Observability` context manager for enabling observability and configuring which Observability Driver to use.
+- `@observable` decorator for selecting which functions/methods to provide observability for.
 
 ### Changed
 - **BREAKING**: `BaseVectorStoreDriver.upsert_text_artifacts` optional arguments are now keyword-only arguments.

diff --git a/docs/griptape-framework/drivers/observability-drivers.md b/docs/griptape-framework/drivers/observability-drivers.md
@@ -0,0 +1,185 @@
+# Observability Drivers
+
+Observability Drivers are used by [Observability](../structures/observability.md) to send telemetry (metrics and traces) related to the execution of an LLM application. The telemetry can be used to monitor the application and to diagnose and troubleshoot issues. All Observability Drivers implement the following methods:
+
+* `__enter__()` sets up the Driver.
+* `__exit__()` tears down the Driver.
+* `observe()` wraps all functions and methods marked with the `@observable` decorator. At a bare minimum, implementations call the wrapped function and return its result (a no-op). This enables the Driver to generate telemetry related to the invocation's call arguments, return values, exceptions, latency, etc.
+
+## Griptape Cloud
+
+The Griptape Cloud Observability Driver instruments `@observable` functions and methods with metrics and traces for use with the Griptape Cloud.
+
+!!! note
+    For the Griptape Cloud Observability Driver to function as intended, it must be run from within either a Managed Structure on Griptape Cloud,
+    or locally via the [Skatepark Emulator](https://github.com/griptape-ai/griptape-cli?tab=readme-ov-file#skatepark-emulator).
+
+Here is an example of how to use the `GriptapeCloudObservabilityDriver` with the `Observability` context manager to send the telemetry to Griptape Cloud:
+
+
+```python title="PYTEST_IGNORE"
+from griptape.drivers import GriptapeCloudObservabilityDriver
+from griptape.rules import Rule
+from griptape.structures import Agent
+from griptape.observability import Observability
+
+observability_driver = GriptapeCloudObservabilityDriver()
+
+with Observability(observability_driver=observability_driver):
+    agent = Agent(rules=[Rule("Output one word")])
+    agent.run("Name an animal")
+```
+
+
+## OpenTelemetry
+
+The [OpenTelemetry](https://opentelemetry.io/) Observability Driver instruments `@observable` functions and methods with metrics and traces for use with OpenTelemetry. You must configure a destination for the telemetry by providing a `SpanProcessor` to the Driver.
+
+Here is an example of how to use the `OpenTelemetryObservabilityDriver` with the `Observability` context manager to output the telemetry directly to the console:
+
+```python title="PYTEST_IGNORE"
+from griptape.drivers import OpenTelemetryObservabilityDriver
+from griptape.rules import Rule
+from griptape.structures import Agent
+from griptape.observability import Observability
+from opentelemetry.sdk.trace.export import ConsoleSpanExporter, BatchSpanProcessor
+
+observability_driver = OpenTelemetryObservabilityDriver(
+    service_name="my-gt-app",
+    span_processor=BatchSpanProcessor(ConsoleSpanExporter())
+)
+
+with Observability(observability_driver=observability_driver):
+    agent = Agent(rules=[Rule("Output one word")])
+    agent.run("Name an animal")
+```
+
+Output (only relevant because of use of `ConsoleSpanExporter`):
+```
+[06/18/24 06:57:22] INFO     PromptTask 2d8ef95bf817480188ae2f74e754308a
+                             Input: Name an animal
+[06/18/24 06:57:23] INFO     PromptTask 2d8ef95bf817480188ae2f74e754308a
+                             Output: Elephant
+{
+    "name": "Agent.before_run()",
+    "context": {
+        "trace_id": "0x4f3d72f7ff4e6a453f5c950fa097583e",
+        "span_id": "0x8cf827b375f6922f",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.INTERNAL",
+    "parent_id": "0x580276d16c584de3",
+    "start_time": "2024-06-18T13:57:22.640040Z",
+    "end_time": "2024-06-18T13:57:22.640822Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {},
+    "events": [],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "my-gt-app"
+        },
+        "schema_url": ""
+    }
+}
+{
+    "name": "Agent.try_run()",
+    "context": {
+        "trace_id": "0x4f3d72f7ff4e6a453f5c950fa097583e",
+        "span_id": "0x7191a27da608cbe7",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.INTERNAL",
+    "parent_id": "0x580276d16c584de3",
+    "start_time": "2024-06-18T13:57:22.640846Z",
+    "end_time": "2024-06-18T13:57:23.287311Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {},
+    "events": [],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "my-gt-app"
+        },
+        "schema_url": ""
+    }
+}
+{
+    "name": "Agent.after_run()",
+    "context": {
+        "trace_id": "0x4f3d72f7ff4e6a453f5c950fa097583e",
+        "span_id": "0x99824dd1bc842f66",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.INTERNAL",
+    "parent_id": "0x580276d16c584de3",
+    "start_time": "2024-06-18T13:57:23.287707Z",
+    "end_time": "2024-06-18T13:57:23.288666Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {},
+    "events": [],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "my-gt-app"
+        },
+        "schema_url": ""
+    }
+}
+{
+    "name": "Agent.run()",
+    "context": {
+        "trace_id": "0x4f3d72f7ff4e6a453f5c950fa097583e",
+        "span_id": "0x580276d16c584de3",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.INTERNAL",
+    "parent_id": "0xa42d36d9fff76325",
+    "start_time": "2024-06-18T13:57:22.640021Z",
+    "end_time": "2024-06-18T13:57:23.288694Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {},
+    "events": [],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "my-gt-app"
+        },
+        "schema_url": ""
+    }
+}
+{
+    "name": "main",
+    "context": {
+        "trace_id": "0x4f3d72f7ff4e6a453f5c950fa097583e",
+        "span_id": "0xa42d36d9fff76325",
+        "trace_state": "[]"
+    },
+    "kind": "SpanKind.INTERNAL",
+    "parent_id": null,
+    "start_time": "2024-06-18T13:57:22.607005Z",
+    "end_time": "2024-06-18T13:57:23.288764Z",
+    "status": {
+        "status_code": "OK"
+    },
+    "attributes": {},
+    "events": [],
+    "links": [],
+    "resource": {
+        "attributes": {
+            "service.name": "my-gt-app"
+        },
+        "schema_url": ""
+    }
+}
+```
+
+
diff --git a/docs/griptape-framework/structures/observability.md b/docs/griptape-framework/structures/observability.md
@@ -0,0 +1,57 @@
+## Overview
+
+The [Observability](../../reference/griptape/observability/observability.md) context manager sends telemetry (metrics and traces) for all functions and methods annotated with the `@observable` decorator to a destination of your choice. This is useful for monitoring and debugging your application.
+
+Observability is completely optional. To opt in, wrap your application code with the [Observability](../../reference/griptape/observability/observability.md) context manager, for example:
+
+```python title="PYTEST_IGNORE"
+from griptape.drivers import GriptapeCloudObservabilityDriver
+from griptape.structures import Agent
+from griptape.observability import Observability
+
+observability_driver = GriptapeCloudObservabilityDriver()
+
+with Observability(observability_driver=observability_driver):
+    # Important! Only code within this block is subject to observability
+    agent = Agent()
+    agent.run("Name the five greatest rappers of all time")
+```
+
+!!! info
+    For available Drivers (and destinations), see [Observability Drivers](../drivers/observability-drivers.md).
+
+## Tracing Custom Code
+
+All functions and methods annotated with the `@observable` decorator will be traced when invoked within the context of the [Observability](../../reference/griptape/observability/observability.md) context manager, including functions and methods defined outside of the Griptape framework. Thus to trace custom code, you just need to add the `@observable` decorator to your function or method, then invoke it within the [Observability](../../reference/griptape/observability/observability.md) context manager.
+
+For example:
+
+```python title="PYTEST_IGNORE"
+import time
+from griptape.drivers import GriptapeCloudObservabilityDriver
+from griptape.rules import Rule
+from griptape.structures import Agent
+from griptape.observability import Observability
+from griptape.common import observable
+
+# Decorate a function
+@observable
+def my_function():
+    time.sleep(3)
+
+class MyClass:
+    # Decorate a method
+    @observable
+    def my_method(self):
+        time.sleep(1)
+        my_function()
+        time.sleep(2)
+
+observability_driver = GriptapeCloudObservabilityDriver()
+
+# When invoking the instrumented code from within the Observability context manager, the
+# telemetry for the custom code will be sent to the destination specified by the driver.
+with Observability(observability_driver=observability_driver):
+    my_function()
+    MyClass().my_method()
+```
diff --git a/griptape/common/__init__.py b/griptape/common/__init__.py
@@ -18,6 +18,7 @@
 
 from .reference import Reference
 
+from .observable import observable, Observable
 
 __all__ = [
     "BaseMessage",
@@ -35,4 +36,6 @@
     "Reference",
     "BaseAction",
     "ToolAction",
+    "observable",
+    "Observable",
 ]
diff --git a/griptape/common/observable.py b/griptape/common/observable.py
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import functools
+from inspect import isfunction
+from typing import Any, Callable, Optional, TypeVar, cast
+
+from attrs import Factory, define, field
+
+T = TypeVar("T", bound=Callable)
+
+
+def observable(*args: T | Any, **kwargs: Any) -> T:
+    return cast(T, Observable(*args, **kwargs))
+
+
+class Observable:
+    @define
+    class Call:
+        func: Callable = field(kw_only=True)
+        instance: Optional[Any] = field(default=None, kw_only=True)
+        args: tuple[Any, ...] = field(default=Factory(tuple), kw_only=True)
+        kwargs: dict[str, Any] = field(default=Factory(dict), kw_only=True)
+        decorator_args: tuple[Any, ...] = field(default=Factory(tuple), kw_only=True)
+        decorator_kwargs: dict[str, Any] = field(default=Factory(dict), kw_only=True)
+
+        def __call__(self) -> Any:
+            # If self.func has a __self__ attribute, it is a bound method and we do not need to pass the instance.
+            args = (self.instance, *self.args) if self.instance and not hasattr(self.func, "__self__") else self.args
+            return self.func(*args, **self.kwargs)
+
+        @property
+        def tags(self) -> Optional[list[str]]:
+            return self.decorator_kwargs.get("tags")
+
+    def __init__(self, *args, **kwargs) -> None:
+        self._instance = None
+        if len(args) == 1 and len(kwargs) == 0 and isfunction(args[0]):
+            # Parameterless call. In otherwords, the `@observable` annotation
+            # was not followed by parentheses.
+            self._func = args[0]
+            functools.update_wrapper(self, self._func)
+            self.decorator_args = ()
+            self.decorator_kwargs = {}
+        else:
+            # Parameterized call. In otherwords, the `@observable` annotation
+            # was followed by parentheses, for example `@observable()`,
+            # `@observable("x")` or `@observable(y="y")`.
+            self._func = None
+            self.decorator_args = args
+            self.decorator_kwargs = kwargs
+
+    def __get__(self, obj: Any, objtype: Any = None) -> Observable:
+        self._instance = obj
+        return self
+
+    def __call__(self, *args, **kwargs) -> Any:
+        if self._func:
+            # Parameterless call (self._func was a set in __init__)
+            from griptape.observability.observability import Observability
+
+            return Observability.observe(
+                Observable.Call(
+                    func=self._func,
+                    instance=self._instance,
+                    args=args,
+                    kwargs=kwargs,
+                    decorator_args=self.decorator_args,
+                    decorator_kwargs=self.decorator_kwargs,
+                )
+            )
+        else:
+            # Parameterized call, create and return the "real" observable decorator
+            func = args[0]
+            decorated_func = Observable(func)
+            decorated_func.decorator_args = self.decorator_args
+            decorated_func.decorator_kwargs = self.decorator_kwargs
+            return decorated_func
diff --git a/griptape/drivers/__init__.py b/griptape/drivers/__init__.py
@@ -109,6 +109,11 @@
 from .audio_transcription.dummy_audio_transcription_driver import DummyAudioTranscriptionDriver
 from .audio_transcription.openai_audio_transcription_driver import OpenAiAudioTranscriptionDriver
 
+from .observability.base_observability_driver import BaseObservabilityDriver
+from .observability.no_op_observability_driver import NoOpObservabilityDriver
+from .observability.open_telemetry_observability_driver import OpenTelemetryObservabilityDriver
+from .observability.griptape_cloud_observability_driver import GriptapeCloudObservabilityDriver
+
 __all__ = [
     "BasePromptDriver",
     "OpenAiChatPromptDriver",
@@ -202,4 +207,8 @@
     "BaseAudioTranscriptionDriver",
     "DummyAudioTranscriptionDriver",
     "OpenAiAudioTranscriptionDriver",
+    "BaseObservabilityDriver",
+    "NoOpObservabilityDriver",
+    "OpenTelemetryObservabilityDriver",
+    "GriptapeCloudObservabilityDriver",
 ]
diff --git a/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py b/griptape/drivers/event_listener/griptape_cloud_event_listener_driver.py
@@ -7,6 +7,7 @@
 from attrs import Attribute, Factory, define, field
 
 from griptape.drivers.event_listener.base_event_listener_driver import BaseEventListenerDriver
+from griptape.events.base_event import BaseEvent
 
 
 @define
@@ -38,6 +39,17 @@ def validate_run_id(self, _: Attribute, structure_run_id: str) -> None:
                 "structure_run_id must be set either in the constructor or as an environment variable (GT_CLOUD_STRUCTURE_RUN_ID).",
             )
 
+    def publish_event(self, event: BaseEvent | dict, *, flush: bool = False) -> None:
+        from griptape.observability.observability import Observability
+
+        event_payload = event.to_dict() if isinstance(event, BaseEvent) else event
+
+        span_id = Observability.get_span_id()
+        if span_id is not None:
+            event_payload["span_id"] = span_id
+
+        super().publish_event(event_payload, flush=flush)
+
     def try_publish_event_payload(self, event_payload: dict) -> None:
         url = urljoin(self.base_url.strip("/"), f"/api/structure-runs/{self.structure_run_id}/events")
 

diff --git a/griptape/drivers/observability/__init__.py b/griptape/drivers/observability/__init__.py