Merge pull request #798 from NVIDIA/feat/observability-support

Feat/observability support

docs/user_guides/configuration-guide.md

@@ -773,6 +773,165 @@ When the `self check input` rail is triggered, the following exception is return
 }
 ```
 
+## Tracing
+
+NeMo Guardrails includes a tracing feature that allows you to monitor and log interactions for better observability and debugging. Tracing can be easily configured via the existing `config.yml` file. Below are the steps to enable and configure tracing in your project.
+
+### Enabling Tracing
+
+To enable tracing, set the enabled flag to true under the tracing section in your `config.yml`:
+
+```yaml
+tracing:
+  enabled: true
+```
+> **Note**: You must install the necessary dependencies to use tracing adapters.
+
+  ```bash
+  pip install "opentelemetry-api opentelemetry-sdk aiofiles"
+  ```
+
+### Configuring Tracing Adapters
+
+Tracing supports multiple adapters that determine how and where the interaction logs are exported. You can configure one or more adapters by specifying them under the adapters list. Below are examples of configuring the built-in `OpenTelemetry` and `FileSystem` adapters:
+
+```yaml
+tracing:
+  enabled: true
+  adapters:
+    - name: OpenTelemetry
+      service_name: "nemo_guardrails_service"
+      exporter: "console"  # Options: "console", "zipkin", etc.
+      resource_attributes:
+        env: "production"
+    - name: FileSystem
+      filepath: './traces/traces.jsonl'
+```
+
+#### OpenTelemetry Adapter
+
+The `OpenTelemetry` adapter integrates with the OpenTelemetry framework, allowing you to export traces to various backends. Key configuration options include:
+
+ • service_name: The name of your service.
+ • exporter: The type of exporter to use (e.g., console, zipkin).
+ • resource_attributes: Additional attributes to include in the trace resource (e.g., environment).
+
+#### FileSystem Adapter
+
+The  `FileSystem`  adapter exports interaction logs to a local JSON Lines file. Key configuration options include:
+
+ • filepath: The path to the file where traces will be stored. If not specified, it defaults to `./.traces/trace.jsonl`.
+
+#### Example Configuration
+
+Here is a complete example of a config.yml with both OpenTelemetry and FileSystem adapters enabled:
+
+```yaml
+tracing:
+  enabled: true
+  adapters:
+    - name: OpenTelemetry
+      service_name: "nemo_guardrails_service"
+      exporter: "console"
+      resource_attributes:
+        env: "production"
+    - name: FileSystem
+      filepath: './traces/traces.jsonl'
+```
+
+### Custom InteractionLogAdapters
+
+NeMo Guardrails allows you to extend its tracing capabilities by creating custom `InteractionLogAdapter` classes. This flexibility enables you to transform and export interaction logs to any backend or format that suits your needs.
+
+#### Implementing a Custom Adapter
+
+To create a custom adapter, you need to implement the `InteractionLogAdapter` abstract base class. Below is the interface you must follow:
+
+```python
+from abc import ABC, abstractmethod
+from nemoguardrails.tracing import InteractionLog
+
+class InteractionLogAdapter(ABC):
+    name: Optional[str] = None
+
+
+    @abstractmethod
+    async def transform_async(self, interaction_log: InteractionLog):
+        """Transforms the InteractionLog into the backend-specific format asynchronously."""
+        raise NotImplementedError
+
+    async def close(self):
+        """Placeholder for any cleanup actions if needed."""
+        pass
+
+    async def __aenter__(self):
+        """Enter the runtime context related to this object."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """Exit the runtime context related to this object."""
+        await self.close()
+
+```
+
+#### Registering Your Custom Adapter
+
+After implementing your custom adapter, you need to register it so that NemoGuardrails can recognize and utilize it. This is done by adding a registration call in your `config.py:`
+
+```python
+from nemoguardrails.tracing.adapters.registry import register_log_adapter
+from path.to.your.adapter import YourCustomAdapter
+
+register_log_adapter(YourCustomAdapter, "CustomLogAdapter")
+```
+
+#### Example: Creating a Custom Adapter
+
+Here’s a simple example of a custom adapter that logs interaction logs to a custom backend:
+
+```python
+from nemoguardrails.tracing.adapters.base import InteractionLogAdapter
+from nemoguardrails.tracing import InteractionLog
+
+class MyCustomLogAdapter(InteractionLogAdapter):
+    name = "MyCustomLogAdapter"
+
+    def __init__(self, custom_option1: str, custom_option2: str):
+      self.custom_option1 = custom_option1
+      self.custom_option2 = custom
+
+    def transform(self, interaction_log: InteractionLog):
+        # Implement your transformation logic here
+        custom_format = convert_to_custom_format(interaction_log)
+        send_to_custom_backend(custom_format)
+
+    async def transform_async(self, interaction_log: InteractionLog):
+        # Implement your asynchronous transformation logic here
+        custom_format = convert_to_custom_format(interaction_log)
+        await send_to_custom_backend_async(custom_format)
+
+    async def close(self):
+        # Implement any necessary cleanup here
+        await cleanup_custom_resources()
+
+```
+
+Updating `config.yml` with Your `CustomLogAdapter`
+
+Once registered, you can configure your custom adapter in the `config.yml` like any other adapter:
+
+```yaml
+tracing:
+  enabled: true
+  adapters:
+    - name: MyCustomLogAdapter
+      custom_option1: "value1"
+      custom_option2: "value2"
+
+```
+
+By following these steps, you can leverage the built-in tracing adapters or create and integrate your own custom adapters to enhance the observability of your NeMo Guardrails powered applications. Whether you choose to export logs to the filesystem, integrate with OpenTelemetry, or implement a bespoke logging solution, tracing provides the flexibility to meet your requirements.
+
 ## Knowledge base Documents
 
 By default, an `LLMRails` instance supports using a set of documents as context for generating the bot responses. To include documents as part of your knowledge base, you must place them in the `kb` folder inside your config folder:

examples/configs/tracing/README.md

@@ -0,0 +1,9 @@
+# README
+
+We encourage you to implement a log adapter for the production environment based on your specific requirements.
+
+To use the `FileSystem` and `OpenTelemetry` adapters, please install the following dependencies:
+
+```bash
+pip install opentelemetry-api opentelemetry-sdk aiofiles
+```

examples/configs/tracing/config.yml

@@ -0,0 +1,15 @@
+models:
+  - type: main
+    engine: openai
+    model: gpt-3.5-turbo-instruct
+
+tracing:
+  enabled: true
+  adapters:
+    - name: OpenTelemetry
+      service_name: "nemo_guardrails_service"
+      exporter: "console"  # Options: "console", "zipkin", etc.
+      resource_attributes:
+        env: "production"
+    - name: FileSystem
+      filepath: './traces/traces.jsonl'

nemoguardrails/rails/llm/config.py

@@ -21,7 +21,7 @@
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import yaml
-from pydantic import BaseModel, ValidationError, root_validator
+from pydantic import BaseModel, ConfigDict, ValidationError, root_validator
 from pydantic.fields import Field
 
 from nemoguardrails import utils
@@ -184,6 +184,19 @@ def check_fields(cls, values):
         return values
 
 
+class LogAdapterConfig(BaseModel):
+    name: str = Field(default="FileSystem", description="The name of the adapter.")
+    model_config = ConfigDict(extra="allow")
+
+
+class TracingConfig(BaseModel):
+    enabled: bool = False
+    adapters: List[LogAdapterConfig] = Field(
+        default_factory=lambda: [LogAdapterConfig()],
+        description="The list of tracing adapters to use. If not specified, the default adapters are used.",
+    )
+
+
 class EmbeddingsCacheConfig(BaseModel):
     """Configuration for the caching embeddings."""
 
@@ -504,6 +517,7 @@ def _join_config(dest_config: dict, additional_config: dict):
         "passthrough",
         "raw_llm_call_action",
         "enable_rails_exceptions",
+        "tracing",
     ]
 
     for field in additional_fields:
@@ -849,6 +863,11 @@ class RailsConfig(BaseModel):
         "This means it will not be altered in any way. ",
     )
 
+    tracing: TracingConfig = Field(
+        default_factory=TracingConfig,
+        description="Configuration for tracing.",
+    )
+
     @root_validator(pre=True, allow_reuse=True)
     def check_prompt_exist_for_self_check_rails(cls, values):
         rails = values.get("rails", {})

nemoguardrails/rails/llm/llmrails.py

@@ -112,6 +112,12 @@ def __init__(
         # Weather the main LLM supports streaming
         self.main_llm_supports_streaming = False
 
+        # InteractionLogAdapters used for tracing
+        if config.tracing:
+            from nemoguardrails.tracing import create_log_adapters
+
+            self._log_adapters = create_log_adapters(config.tracing)
+
         # We also load the default flows from the `default_flows.yml` file in the current folder.
         # But only for version 1.0.
         # TODO: decide on the default flows for 2.x.
@@ -789,6 +795,19 @@ async def generate_async(
             # print("Closing the stream handler explicitly")
             await streaming_handler.push_chunk(None)
 
+        # IF tracing is enabled we need to set GenerationLog attrs
+        if self.config.tracing.enabled:
+            if options is None:
+                options = GenerationOptions()
+            if (
+                not options.log.activated_rails
+                or not options.log.llm_calls
+                or not options.log.internal_events
+            ):
+                options.log.activated_rails = True
+                options.log.llm_calls = True
+                options.log.internal_events = True
+
         # If we have generation options, we prepare a GenerationResponse instance.
         if options:
             # If a prompt was used, we only need to return the content of the message.
@@ -881,6 +900,17 @@ async def generate_async(
             if state is not None:
                 res.state = output_state
 
+            if self.config.tracing.enabled:
+                # TODO: move it to the top once resolved circular dependency of eval
+                # lazy import to avoid circular dependency
+                from nemoguardrails.tracing import Tracer
+
+                # Create a Tracer instance with instantiated adapters
+                tracer = Tracer(
+                    input=messages, response=res, adapters=self._log_adapters
+                )
+                await tracer.export_async()
+                res = res.response[0]
             return res
         else:
             # If a prompt is used, we only return the content of the message.

nemoguardrails/tracing/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .tracer import InteractionLog, Tracer, create_log_adapters

nemoguardrails/tracing/adapters/__init__.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from .filesystem import FileSystemAdapter
+from .registry import register_log_adapter
+
+register_log_adapter(FileSystemAdapter, "FileSystem")
+
+try:
+    from .opentelemetry import OpenTelemetryAdapter
+
+    register_log_adapter(OpenTelemetryAdapter, "OpenTelemetry")
+
+except ImportError:
+    pass
+
+# __all__ = ["InteractionLogAdapter", "LogAdapterRegistry"]

nemoguardrails/tracing/adapters/base.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from nemoguardrails.eval.models import InteractionLog
+
+
+class InteractionLogAdapter(ABC):
+    name: Optional[str] = None
+
+    @abstractmethod
+    def transform(self, interaction_log: InteractionLog):
+        """Transforms the InteractionLog into the backend-specific format."""
+        pass
+
+    @abstractmethod
+    async def transform_async(self, interaction_log: InteractionLog):
+        """Transforms the InteractionLog into the backend-specific format asynchronously."""
+        raise NotImplementedError
+
+    async def close(self):
+        """Placeholder for any cleanup actions if needed."""
+        pass
+
+    async def __aenter__(self):
+        """Enter the runtime context related to this object."""
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        """Exit the runtime context related to this object."""
+        await self.close()