From 14a6ce367d2d23057d8131dbecba127b3c345b69 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 15:40:58 -0700 Subject: [PATCH] add types for BedrockMessage --- litellm/types/guardrails.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 66c2a535ad6..13992beec58 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Dict, List, Optional, TypedDict +from typing import Dict, List, Literal, Optional, TypedDict from pydantic import BaseModel, ConfigDict from typing_extensions import Required, TypedDict @@ -76,8 +76,14 @@ class LitellmParams(TypedDict, total=False): mode: str api_key: str api_base: Optional[str] + + # Lakera specific params category_thresholds: Optional[LakeraCategoryThresholds] + # Bedrock specific params + guardrailIdentifier: Optional[str] + guardrailVersion: Optional[str] + class Guardrail(TypedDict): guardrail_name: str @@ -92,3 +98,16 @@ class GuardrailEventHooks(str, Enum): pre_call = "pre_call" post_call = "post_call" during_call = "during_call" + + +class BedrockTextContent(TypedDict): + text: str + + +class BedrockContentItem(TypedDict): + text: BedrockTextContent + + +class BedrockMessage(TypedDict): + source: Literal["INPUT", "OUTPUT"] + content: List[BedrockContentItem]