AI Design Pattern: Guardrails

TOPIC_GUARDRAIL_PROMPT = """
Assess whether the user question is allowed or not. The allowed topics are cats and dogs.
"""

class TopicGuardrailResult(BaseModel):
    allowed: bool


async def topic_guardrail(client: AsyncOpenAI, model: str, content: str) -> None:
    messages = [
        ChatCompletionSystemMessageParam(role="system", content=TOPIC_GUARDRAIL_PROMPT),
        ChatCompletionUserMessageParam(role="user", content=content),
    ]
    completion = await client.beta.chat.completions.parse(
        model=model,
        messages=messages,
        temperature=0.0,
        response_format=TopicGuardrailResult,
    )
    parsed = completion.choices[0].message.parsed

    if parsed is None or not parsed.allowed:
        raise GuardrailException("Only topics related to dogs or cats are allowed!")

DOMAIN = "animal breed recommendation"

ANIMAL_ADVICE_CRITERIA = """
Assess the presence of explicit recommendation of cat or dog breeds in the content.
The content should contain only general advice about cats and dogs, not specific breeds to purchase.
"""

ANIMAL_ADVICE_STEPS = """
1. Read the content and the criteria carefully.
2. Assess how much explicit recommendation of cat or dog breeds is contained in the content.
3. Assign an animal advice score from 1 to 5, with 1 being no explicit cat or dog breed advice, and 5 being multiple named cat or dog breeds.
"""

MODERATION_GUARDRAIL_PROMPT = """
You are a moderation assistant. Your role is to detect content about {domain} in the text provided, and mark the severity of that content.

## {domain}

### Criteria

{scoring_criteria}

### Instructions

{scoring_steps}

### Content

{content}

### Evaluation (score only!)
"""

class ModerationGuardrailResult(BaseModel):
    score: int


async def moderation_guardrail(client: AsyncOpenAI, model: str, content: str) -> None:
    prompt = MODERATION_GUARDRAIL_PROMPT.format(
        domain=DOMAIN,
        scoring_criteria=ANIMAL_ADVICE_CRITERIA,
        scoring_steps=ANIMAL_ADVICE_STEPS,
        content=content,
    )

    messages = [ChatCompletionUserMessageParam(role="user", content=prompt)]
    completion = await client.beta.chat.completions.parse(
        model=model,
        messages=messages,
        temperature=0.0,
        response_format=ModerationGuardrailResult,
    )
    parsed = completion.choices[0].message.parsed

    if parsed is None or parsed.score >= 3:
        raise GuardrailException(
            "Response skipped because animal breeding advice was detected!"
        )

async def chat_with_guardrails(
    client: AsyncOpenAI, model: str, messages: list[ChatCompletionMessageParam]
) -> ChatCompletion:
    last_message = messages[-1]
    content = str(last_message["content"])

    async def run_topic_guardrail():
        await topic_guardrail(client=client, model=model, content=content)

    async def run_completion_and_moderation_guardrail():
        completion = await client.chat.completions.create(
            model=model, messages=messages
        )
        content = completion.choices[0].message.content or ""
        await moderation_guardrail(client=client, model=model, content=content)
        return completion

    try:
        topic_task = asyncio.create_task(run_topic_guardrail())
        completion_task = asyncio.create_task(run_completion_and_moderation_guardrail())

        # Wait for topic check to complete
        await topic_task

        # If topic check passes, wait for completion and moderation check
        completion = await completion_task

        return completion

    except GuardrailException as e:
        completion_task.cancel()
        error_message = str(e)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")

    # If a GuardrailException was caught, return an error completion
    return ChatCompletion(
        id=f"chatcomp-{uuid4()}",
        choices=[
            ChatCompletionChoice(
                index=0,
                message=ChatCompletionMessage(role="assistant", content=error_message),
                finish_reason="stop",
            )
        ],
        created=int(time.time()),
        model=model,
        object="chat.completion",
    )

class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[ChatCompletionMessageParam]


def create_fastapi_app(client: AsyncOpenAI) -> FastAPI:
    fastapi_app = FastAPI()

    @fastapi_app.post("/v1/chat/completions", response_model=ChatCompletion)
    async def chat(request: ChatCompletionRequest) -> ChatCompletion:
        return await chat_with_guardrails(
            client=client, model=request.model, messages=request.messages
        )

    return fastapi_app

curl -X POST "http://0.0.0.0:8000/v1/chat/completions" \
     -H "Content-Type: application/json" \
     -d '{"model":"gpt-4o-2024-08-06","messages":[{"role":"user","content":"How can I introduce a new dog to my cat?"}]}' | jq

from openai import OpenAI

client = OpenAI(base_url="http://0.0.0.0:8000/v1")

completion = client.chat.completions.create(
    model="gpt-4o-2024-08-06",
    messages=[{"role": "user", "content": "How can I introduce a new dog to my cat?"}],
)

from langfuse.decorators import observe


@observe()
async def topic_guardrail(client: AsyncOpenAI, model: str, content: str) -> None:
  ...

@observe()
async def moderation_guardrail(client: AsyncOpenAI, model: str, content: str) -> None:
  ...

@observe()
async def chat_with_guardrails(
    client: AsyncOpenAI, model: str, messages: list[ChatCompletionMessageParam]
) -> ChatCompletion:
  ...