From 4aa1fe62d56c418c51f74a29b441d7be0daf2554 Mon Sep 17 00:00:00 2001 From: Konrad Czarnota Date: Mon, 4 Nov 2024 14:25:07 +0100 Subject: [PATCH 1/2] enabler: guardrails component design --- examples/core/guardrail.py | 29 ++++++++++ .../src/ragbits/core/guardrails/__init__.py | 0 .../src/ragbits/core/guardrails/base.py | 54 +++++++++++++++++++ .../core/guardrails/openai_moderation.py | 51 ++++++++++++++++++ .../tests/unit/guardrails/__init__.py | 0 .../unit/guardrails/test_openai_moderation.py | 53 ++++++++++++++++++ .../providers/unstructured/images.py | 5 +- 7 files changed, 188 insertions(+), 4 deletions(-) create mode 100644 examples/core/guardrail.py create mode 100644 packages/ragbits-core/src/ragbits/core/guardrails/__init__.py create mode 100644 packages/ragbits-core/src/ragbits/core/guardrails/base.py create mode 100644 packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py create mode 100644 packages/ragbits-core/tests/unit/guardrails/__init__.py create mode 100644 packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py diff --git a/examples/core/guardrail.py b/examples/core/guardrail.py new file mode 100644 index 000000000..13b976750 --- /dev/null +++ b/examples/core/guardrail.py @@ -0,0 +1,29 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "ragbits-core", +# "openai", +# ] +# /// +import asyncio +from argparse import ArgumentParser + +from ragbits.core.guardrails.base import GuardrailManager +from ragbits.core.guardrails.openai_moderation import OpenAIModerationGuardrail + + +async def guardrail_run(message: str) -> None: + """ + Example of using the OpenAIModerationGuardrail. Requires the OPENAI_API_KEY environment variable to be set. + """ + manager = GuardrailManager([OpenAIModerationGuardrail()]) + res = await manager.verify(message) + print(res) + + +if __name__ == "__main__": + args = ArgumentParser() + args.add_argument("message", nargs="+", type=str, help="Message to validate") + parsed_args = args.parse_args() + + asyncio.run(guardrail_run("".join(parsed_args.message))) diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/__init__.py b/packages/ragbits-core/src/ragbits/core/guardrails/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/base.py b/packages/ragbits-core/src/ragbits/core/guardrails/base.py new file mode 100644 index 000000000..ccd12d251 --- /dev/null +++ b/packages/ragbits-core/src/ragbits/core/guardrails/base.py @@ -0,0 +1,54 @@ +from abc import ABC, abstractmethod + +from pydantic import BaseModel + +from ragbits.core.prompt import Prompt + + +class GuardrailVerificationResult(BaseModel): + """ + Class representing result of guardrail verification + """ + + guardrail_name: str + succeeded: bool + fail_reason: str | None + + +class Guardrail(ABC): + """ + Abstract class representing guardrail + """ + + @abstractmethod + async def verify(self, input_to_verify: Prompt | str) -> GuardrailVerificationResult: + """ + Verifies whether provided input meets certain criteria + + Args: + input_to_verify: prompt or output of the model to check + + Returns: + verification result + """ + + +class GuardrailManager: + """ + Class responsible for running guardrails + """ + + def __init__(self, guardrails: list[Guardrail]): + self._guardrails = guardrails + + async def verify(self, input_to_verify: Prompt | str) -> list[GuardrailVerificationResult]: + """ + Verifies whether provided input meets certain criteria + + Args: + input_to_verify: prompt or output of the model to check + + Returns: + list of verification result + """ + return [await guardrail.verify(input_to_verify) for guardrail in self._guardrails] diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py b/packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py new file mode 100644 index 000000000..4e0fd6109 --- /dev/null +++ b/packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py @@ -0,0 +1,51 @@ +import base64 + +from openai import AsyncOpenAI + +from ragbits.core.guardrails.base import Guardrail, GuardrailVerificationResult +from ragbits.core.prompt import Prompt + + +class OpenAIModerationGuardrail(Guardrail): + """ + Guardrail based on OpenAI moderation + """ + + def __init__(self, moderation_model: str = "omni-moderation-latest"): + self._openai_client = AsyncOpenAI() + self._moderation_model = moderation_model + + async def verify(self, input_to_verify: Prompt | str) -> GuardrailVerificationResult: + """ + Verifies whether provided input meets certain criteria + + Args: + input_to_verify: prompt or output of the model to check + + Returns: + verification result + """ + if isinstance(input_to_verify, Prompt): + inputs = [{"type": "text", "text": input_to_verify.rendered_user_prompt}] + if input_to_verify.rendered_system_prompt is not None: + inputs.append({"type": "text", "text": input_to_verify.rendered_system_prompt}) + if images := input_to_verify.images: + inputs.extend( + [ + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{base64.b64encode(im).decode('utf-8')}"}, # type: ignore + } + for im in images + ] + ) + else: + inputs = [{"type": "text", "text": input_to_verify}] + response = await self._openai_client.moderations.create(model=self._moderation_model, input=inputs) # type: ignore + + fail_reasons = [result for result in response.results if result.flagged] + return GuardrailVerificationResult( + guardrail_name=self.__class__.__name__, + succeeded=len(fail_reasons) == 0, + fail_reason=None if len(fail_reasons) == 0 else str(fail_reasons), + ) diff --git a/packages/ragbits-core/tests/unit/guardrails/__init__.py b/packages/ragbits-core/tests/unit/guardrails/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py b/packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py new file mode 100644 index 000000000..599f1067b --- /dev/null +++ b/packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py @@ -0,0 +1,53 @@ +import os +from unittest.mock import AsyncMock, patch + +from pydantic import BaseModel + +from ragbits.core.guardrails.base import GuardrailManager, GuardrailVerificationResult +from ragbits.core.guardrails.openai_moderation import OpenAIModerationGuardrail + + +class MockedModeration(BaseModel): + flagged: bool + fail_reason: str | None + + +class MockedModerationCreateResponse(BaseModel): + results: list[MockedModeration] + + +async def test_manager(): + guardrail_mock = AsyncMock() + guardrail_mock.verify.return_value = GuardrailVerificationResult( + guardrail_name=".", succeeded=True, fail_reason=None + ) + manager = GuardrailManager([guardrail_mock]) + results = await manager.verify("test") + assert guardrail_mock.verify.call_count == 1 + assert len(results) == 1 + + +@patch.dict(os.environ, {"OPENAI_API_KEY": "."}, clear=True) +async def test_not_flagged(): + guardrail = OpenAIModerationGuardrail() + guardrail._openai_client = AsyncMock() + guardrail._openai_client.moderations.create.return_value = MockedModerationCreateResponse( + results=[MockedModeration(flagged=False, fail_reason=None)] + ) + results = await guardrail.verify("Test") + assert results.succeeded is True + assert results.fail_reason is None + assert results.guardrail_name == "OpenAIModerationGuardrail" + + +@patch.dict(os.environ, {"OPENAI_API_KEY": "."}, clear=True) +async def test_flagged(): + guardrail = OpenAIModerationGuardrail() + guardrail._openai_client = AsyncMock() + guardrail._openai_client.moderations.create.return_value = MockedModerationCreateResponse( + results=[MockedModeration(flagged=True, fail_reason="Harmful content")] + ) + results = await guardrail.verify("Test") + assert results.succeeded is False + assert results.fail_reason == "[MockedModeration(flagged=True, fail_reason='Harmful content')]" + assert results.guardrail_name == "OpenAIModerationGuardrail" diff --git a/packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py b/packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py index 4c2280a5b..674a1584f 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py @@ -21,8 +21,8 @@ to_text_element, ) -DEFAULT_LLM_IMAGE_SUMMARIZATION_MODEL = "gpt-4o-mini" DEFAULT_IMAGE_QUESTION_PROMPT = "Describe the content of the image." +DEFAULT_LLM_IMAGE_DESCRIPTION_MODEL = "gpt-4o-mini" class _ImagePrompt(Prompt): @@ -34,9 +34,6 @@ class _ImagePromptInput(BaseModel): images: list[bytes] -DEFAULT_LLM_IMAGE_DESCRIPTION_MODEL = "gpt-4o-mini" - - class UnstructuredImageProvider(UnstructuredDefaultProvider): """ A specialized provider that handles pngs and jpgs using the Unstructured From a566c4f12347733a915de67b5513b45e6bcfd54a Mon Sep 17 00:00:00 2001 From: Konrad Czarnota Date: Tue, 5 Nov 2024 08:55:25 +0100 Subject: [PATCH 2/2] Move to separate package --- .../openai_moderation.py} | 4 +- .../CHANGELOG.md} | 0 packages/ragbits-guardrails/README.md | 1 + .../py.typed} | 0 packages/ragbits-guardrails/pyproject.toml | 61 +++++++++++++++++++ .../src/ragbits/guardrails/__init__.py | 0 .../src/ragbits}/guardrails/base.py | 0 .../ragbits}/guardrails/openai_moderation.py | 2 +- .../tests/unit}/test_openai_moderation.py | 4 +- pyproject.toml | 4 ++ uv.lock | 40 ++++++++++++ 11 files changed, 111 insertions(+), 5 deletions(-) rename examples/{core/guardrail.py => guardrails/openai_moderation.py} (83%) rename packages/{ragbits-core/src/ragbits/core/guardrails/__init__.py => ragbits-guardrails/CHANGELOG.md} (100%) create mode 100644 packages/ragbits-guardrails/README.md rename packages/{ragbits-core/tests/unit/guardrails/__init__.py => ragbits-guardrails/py.typed} (100%) create mode 100644 packages/ragbits-guardrails/pyproject.toml create mode 100644 packages/ragbits-guardrails/src/ragbits/guardrails/__init__.py rename packages/{ragbits-core/src/ragbits/core => ragbits-guardrails/src/ragbits}/guardrails/base.py (100%) rename packages/{ragbits-core/src/ragbits/core => ragbits-guardrails/src/ragbits}/guardrails/openai_moderation.py (95%) rename packages/{ragbits-core/tests/unit/guardrails => ragbits-guardrails/tests/unit}/test_openai_moderation.py (91%) diff --git a/examples/core/guardrail.py b/examples/guardrails/openai_moderation.py similarity index 83% rename from examples/core/guardrail.py rename to examples/guardrails/openai_moderation.py index 13b976750..75e42155d 100644 --- a/examples/core/guardrail.py +++ b/examples/guardrails/openai_moderation.py @@ -8,8 +8,8 @@ import asyncio from argparse import ArgumentParser -from ragbits.core.guardrails.base import GuardrailManager -from ragbits.core.guardrails.openai_moderation import OpenAIModerationGuardrail +from ragbits.guardrails.base import GuardrailManager +from ragbits.guardrails.openai_moderation import OpenAIModerationGuardrail async def guardrail_run(message: str) -> None: diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/__init__.py b/packages/ragbits-guardrails/CHANGELOG.md similarity index 100% rename from packages/ragbits-core/src/ragbits/core/guardrails/__init__.py rename to packages/ragbits-guardrails/CHANGELOG.md diff --git a/packages/ragbits-guardrails/README.md b/packages/ragbits-guardrails/README.md new file mode 100644 index 000000000..4dcb2479e --- /dev/null +++ b/packages/ragbits-guardrails/README.md @@ -0,0 +1 @@ +# Ragbits Guardrails diff --git a/packages/ragbits-core/tests/unit/guardrails/__init__.py b/packages/ragbits-guardrails/py.typed similarity index 100% rename from packages/ragbits-core/tests/unit/guardrails/__init__.py rename to packages/ragbits-guardrails/py.typed diff --git a/packages/ragbits-guardrails/pyproject.toml b/packages/ragbits-guardrails/pyproject.toml new file mode 100644 index 000000000..d398e8b06 --- /dev/null +++ b/packages/ragbits-guardrails/pyproject.toml @@ -0,0 +1,61 @@ +[project] +name = "ragbits-guardrails" +version = "0.2.0" +description = "Guardrails module for Ragbits components" +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +authors = [ + { name = "deepsense.ai", email = "ragbits@deepsense.ai"} +] +keywords = [ + "Retrieval Augmented Generation", + "RAG", + "Large Language Models", + "LLMs", + "Generative AI", + "GenAI", + "Evaluation" +] +classifiers = [ + "Development Status :: 4 - Beta", + "Environment :: Console", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", +] +dependencies = ["ragbits-core==0.2.0"] + +[project.optional-dependencies] +openai = [ + "openai~=1.51.0", +] + +[tool.uv] +dev-dependencies = [ + "pre-commit~=3.8.0", + "pytest~=8.3.3", + "pytest-cov~=5.0.0", + "pytest-asyncio~=0.24.0", + "pip-licenses>=4.0.0,<5.0.0" +] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.metadata] +allow-direct-references = true + +[tool.hatch.build.targets.wheel] +packages = ["src/ragbits"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" diff --git a/packages/ragbits-guardrails/src/ragbits/guardrails/__init__.py b/packages/ragbits-guardrails/src/ragbits/guardrails/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/base.py b/packages/ragbits-guardrails/src/ragbits/guardrails/base.py similarity index 100% rename from packages/ragbits-core/src/ragbits/core/guardrails/base.py rename to packages/ragbits-guardrails/src/ragbits/guardrails/base.py diff --git a/packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py b/packages/ragbits-guardrails/src/ragbits/guardrails/openai_moderation.py similarity index 95% rename from packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py rename to packages/ragbits-guardrails/src/ragbits/guardrails/openai_moderation.py index 4e0fd6109..bbc36a2ac 100644 --- a/packages/ragbits-core/src/ragbits/core/guardrails/openai_moderation.py +++ b/packages/ragbits-guardrails/src/ragbits/guardrails/openai_moderation.py @@ -2,8 +2,8 @@ from openai import AsyncOpenAI -from ragbits.core.guardrails.base import Guardrail, GuardrailVerificationResult from ragbits.core.prompt import Prompt +from ragbits.guardrails.base import Guardrail, GuardrailVerificationResult class OpenAIModerationGuardrail(Guardrail): diff --git a/packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py b/packages/ragbits-guardrails/tests/unit/test_openai_moderation.py similarity index 91% rename from packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py rename to packages/ragbits-guardrails/tests/unit/test_openai_moderation.py index 599f1067b..4831946b7 100644 --- a/packages/ragbits-core/tests/unit/guardrails/test_openai_moderation.py +++ b/packages/ragbits-guardrails/tests/unit/test_openai_moderation.py @@ -3,8 +3,8 @@ from pydantic import BaseModel -from ragbits.core.guardrails.base import GuardrailManager, GuardrailVerificationResult -from ragbits.core.guardrails.openai_moderation import OpenAIModerationGuardrail +from ragbits.guardrails.base import GuardrailManager, GuardrailVerificationResult +from ragbits.guardrails.openai_moderation import OpenAIModerationGuardrail class MockedModeration(BaseModel): diff --git a/pyproject.toml b/pyproject.toml index 3e2499ba0..21c7a53f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ dependencies = [ "ragbits-core[litellm,local,lab,chroma]", "ragbits-document-search[gcs, huggingface]", "ragbits-evaluate[relari]", + "ragbits-guardrails[openai]", ] [tool.uv] @@ -35,6 +36,7 @@ ragbits-cli = { workspace = true } ragbits-core = { workspace = true } ragbits-document-search = { workspace = true } ragbits-evaluate = {workspace = true} +ragbits-guardrails = {workspace = true} [tool.uv.workspace] members = [ @@ -42,6 +44,7 @@ members = [ "packages/ragbits-core", "packages/ragbits-document-search", "packages/ragbits-evaluate", + "packages/ragbits-guardrails", ] [tool.pytest] @@ -88,6 +91,7 @@ mypy_path = [ "packages/ragbits-core/src", "packages/ragbits-document-search/src", "packages/ragbits-evaluate/src", + "packages/ragbits-guardrails/src", ] exclude = ["scripts"] diff --git a/uv.lock b/uv.lock index 15bf71910..f20c894d2 100644 --- a/uv.lock +++ b/uv.lock @@ -13,6 +13,7 @@ members = [ "ragbits-core", "ragbits-document-search", "ragbits-evaluate", + "ragbits-guardrails", "ragbits-workspace", ] @@ -3721,6 +3722,43 @@ dev = [ { name = "pytest-cov", specifier = "~=5.0.0" }, ] +[[package]] +name = "ragbits-guardrails" +version = "0.2.0" +source = { editable = "packages/ragbits-guardrails" } +dependencies = [ + { name = "ragbits-core" }, +] + +[package.optional-dependencies] +openai = [ + { name = "openai" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pip-licenses" }, + { name = "pre-commit" }, + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "pytest-cov" }, +] + +[package.metadata] +requires-dist = [ + { name = "openai", marker = "extra == 'openai'", specifier = "~=1.51.0" }, + { name = "ragbits-core", editable = "packages/ragbits-core" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "pip-licenses", specifier = ">=4.0.0,<5.0.0" }, + { name = "pre-commit", specifier = "~=3.8.0" }, + { name = "pytest", specifier = "~=8.3.3" }, + { name = "pytest-asyncio", specifier = "~=0.24.0" }, + { name = "pytest-cov", specifier = "~=5.0.0" }, +] + [[package]] name = "ragbits-workspace" version = "0.1.0" @@ -3730,6 +3768,7 @@ dependencies = [ { name = "ragbits-core", extra = ["chroma", "lab", "litellm", "local"] }, { name = "ragbits-document-search", extra = ["gcs", "huggingface"] }, { name = "ragbits-evaluate", extra = ["relari"] }, + { name = "ragbits-guardrails", extra = ["openai"] }, ] [package.dev-dependencies] @@ -3757,6 +3796,7 @@ requires-dist = [ { name = "ragbits-core", extras = ["litellm", "local", "lab", "chroma"], editable = "packages/ragbits-core" }, { name = "ragbits-document-search", extras = ["gcs", "huggingface"], editable = "packages/ragbits-document-search" }, { name = "ragbits-evaluate", extras = ["relari"], editable = "packages/ragbits-evaluate" }, + { name = "ragbits-guardrails", extras = ["openai"], editable = "packages/ragbits-guardrails" }, ] [package.metadata.requires-dev]