From 7c12cda244b442d60b64cdc4069b954ae3a2348e Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Thu, 26 Dec 2024 18:18:01 -0800 Subject: [PATCH] llama guard --- .../inline/safety/llama_guard/llama_guard.py | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/llama_stack/providers/inline/safety/llama_guard/llama_guard.py b/llama_stack/providers/inline/safety/llama_guard/llama_guard.py index bbdd5c3df..00213ac83 100644 --- a/llama_stack/providers/inline/safety/llama_guard/llama_guard.py +++ b/llama_stack/providers/inline/safety/llama_guard/llama_guard.py @@ -9,10 +9,24 @@ import re from string import Template from typing import Any, Dict, List, Optional -from llama_models.llama3.api.datatypes import * # noqa: F403 -from llama_stack.apis.inference import * # noqa: F403 -from llama_stack.apis.safety import * # noqa: F403 +from llama_models.datatypes import CoreModelId +from llama_models.llama3.api.datatypes import Role + from llama_stack.apis.common.content_types import ImageContentItem, TextContentItem +from llama_stack.apis.inference import ( + ChatCompletionResponseEventType, + Inference, + Message, + UserMessage, +) +from llama_stack.apis.safety import ( + RunShieldResponse, + Safety, + SafetyViolation, + ViolationLevel, +) + +from llama_stack.apis.shields import Shield from llama_stack.distribution.datatypes import Api from llama_stack.providers.datatypes import ShieldsProtocolPrivate