Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat/observability support #798

Merged
merged 29 commits into from
Oct 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
5ee2d25
feat: Implement tracing module
Pouyanpi Oct 10, 2024
c2bf0fd
feat(tracing): add JSON and OpenTelemetry adapters
Pouyanpi Oct 10, 2024
752ce9e
feat(config): add tracing configuration support
Pouyanpi Oct 10, 2024
31bf67b
refactor(tracing): restructure adapters to avoid circular imports
Pouyanpi Oct 10, 2024
471427a
feat(tracing): add TracingConfig support to Tracer
Pouyanpi Oct 10, 2024
25d4665
fix(adapter_factory): pop name from config before passing to class in…
Pouyanpi Oct 11, 2024
7b95bff
refactor: rename JsonAdapter to FileSystemAdapter
Pouyanpi Oct 11, 2024
0bcf05a
feat(tracing): add tracing support to LLMRails
Pouyanpi Oct 11, 2024
4b0bf25
feat(tests): add tests for tracing adapters
Pouyanpi Oct 11, 2024
2ddb6c7
refactor: rename AdapterConfig to LogAdapterConfig
Pouyanpi Oct 14, 2024
65e6cf6
feat: add async export method to Tracer class
Pouyanpi Oct 14, 2024
5d15837
fix: change trace file extension to .jsonl
Pouyanpi Oct 14, 2024
fa2981b
feat: Implement LogAdapterRegistry for tracing adapters
Pouyanpi Oct 14, 2024
f22d3a9
feat: register tracing adapters in __init__.py
Pouyanpi Oct 14, 2024
124097f
feat: add async support to tracing adapters
Pouyanpi Oct 14, 2024
360b35d
feat: add async support and adapter factory to tracer
Pouyanpi Oct 14, 2024
91c66df
feat(tracing): integrate async tracing in LLMRails
Pouyanpi Oct 14, 2024
47a932a
test(tracing): add async tracing tests
Pouyanpi Oct 14, 2024
728956d
chore: remove adapter_factory
Pouyanpi Oct 14, 2024
2aacec8
fix: handle missing OpenTelemetryAdapter gracefully
Pouyanpi Oct 14, 2024
ee517f7
test: skip test if aiofiles is not installed
Pouyanpi Oct 14, 2024
a33602e
docs: add tracing configuration guide
Pouyanpi Oct 14, 2024
2483ad5
fix: reinitialize executor if shut down in OpenTelemetryAdapter
Pouyanpi Oct 16, 2024
cac4b99
feat: add example configuration for tracing
Pouyanpi Oct 16, 2024
a91055a
feat: add tracing dependencies to pyproject.toml
Pouyanpi Oct 16, 2024
ff0c9f2
feat(config): set default values for LogAdapterConfig
Pouyanpi Oct 16, 2024
6927702
refactor: simplify OpenTelemetryAdapter initialization and fix unexpe…
Pouyanpi Oct 16, 2024
13e3f11
refactor: update OpenTelemetryAdapter unit tests
Pouyanpi Oct 16, 2024
4be3307
fix: ensure correct response extraction in LLMRails to conform curren…
Pouyanpi Oct 16, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 159 additions & 0 deletions docs/user_guides/configuration-guide.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
9 changes: 9 additions & 0 deletions examples/configs/tracing/README.md
Original file line number Diff line number Diff line change
@@ -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
```
15 changes: 15 additions & 0 deletions examples/configs/tracing/config.yml
Original file line number Diff line number Diff line change
@@ -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'
21 changes: 20 additions & 1 deletion nemoguardrails/rails/llm/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.colang import parse_colang_file, parse_flow_elements
Expand Down Expand Up @@ -183,6 +183,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."""

Expand Down Expand Up @@ -503,6 +516,7 @@ def _join_config(dest_config: dict, additional_config: dict):
"passthrough",
"raw_llm_call_action",
"enable_rails_exceptions",
"tracing",
]

for field in additional_fields:
Expand Down Expand Up @@ -836,6 +850,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", {})
Expand Down
30 changes: 30 additions & 0 deletions nemoguardrails/rails/llm/llmrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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(
Pouyanpi marked this conversation as resolved.
Show resolved Hide resolved
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.
Expand Down
16 changes: 16 additions & 0 deletions nemoguardrails/tracing/__init__.py
Original file line number Diff line number Diff line change
@@ -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
30 changes: 30 additions & 0 deletions nemoguardrails/tracing/adapters/__init__.py
Original file line number Diff line number Diff line change
@@ -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"]
45 changes: 45 additions & 0 deletions nemoguardrails/tracing/adapters/base.py
Original file line number Diff line number Diff line change
@@ -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()
Loading