llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/impls/meta_reference/memory/faiss.py

@@ -1,113 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-import logging
-
-from typing import Any, Dict, List, Optional
-
-import faiss
-import numpy as np
-from numpy.typing import NDArray
-
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-
-from llama_stack.apis.memory import *  # noqa: F403
-from llama_stack.providers.datatypes import MemoryBanksProtocolPrivate
-
-from llama_stack.providers.utils.memory.vector_store import (
-    ALL_MINILM_L6_V2_DIMENSION,
-    BankWithIndex,
-    EmbeddingIndex,
-)
-from llama_stack.providers.utils.telemetry import tracing
-
-from .config import FaissImplConfig
-
-logger = logging.getLogger(__name__)
-
-
-class FaissIndex(EmbeddingIndex):
-    id_by_index: Dict[int, str]
-    chunk_by_index: Dict[int, str]
-
-    def __init__(self, dimension: int):
-        self.index = faiss.IndexFlatL2(dimension)
-        self.id_by_index = {}
-        self.chunk_by_index = {}
-
-    @tracing.span(name="add_chunks")
-    async def add_chunks(self, chunks: List[Chunk], embeddings: NDArray):
-        indexlen = len(self.id_by_index)
-        for i, chunk in enumerate(chunks):
-            self.chunk_by_index[indexlen + i] = chunk
-            self.id_by_index[indexlen + i] = chunk.document_id
-
-        self.index.add(np.array(embeddings).astype(np.float32))
-
-    async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
-        distances, indices = self.index.search(
-            embedding.reshape(1, -1).astype(np.float32), k
-        )
-
-        chunks = []
-        scores = []
-        for d, i in zip(distances[0], indices[0]):
-            if i < 0:
-                continue
-            chunks.append(self.chunk_by_index[int(i)])
-            scores.append(1.0 / float(d))
-
-        return QueryDocumentsResponse(chunks=chunks, scores=scores)
-
-
-class FaissMemoryImpl(Memory, MemoryBanksProtocolPrivate):
-    def __init__(self, config: FaissImplConfig) -> None:
-        self.config = config
-        self.cache = {}
-
-    async def initialize(self) -> None: ...
-
-    async def shutdown(self) -> None: ...
-
-    async def register_memory_bank(
-        self,
-        memory_bank: MemoryBankDef,
-    ) -> None:
-        assert (
-            memory_bank.type == MemoryBankType.vector.value
-        ), f"Only vector banks are supported {memory_bank.type}"
-
-        index = BankWithIndex(
-            bank=memory_bank, index=FaissIndex(ALL_MINILM_L6_V2_DIMENSION)
-        )
-        self.cache[memory_bank.identifier] = index
-
-    async def list_memory_banks(self) -> List[MemoryBankDef]:
-        return [i.bank for i in self.cache.values()]
-
-    async def insert_documents(
-        self,
-        bank_id: str,
-        documents: List[MemoryBankDocument],
-        ttl_seconds: Optional[int] = None,
-    ) -> None:
-        index = self.cache.get(bank_id)
-        if index is None:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        await index.insert_documents(documents)
-
-    async def query_documents(
-        self,
-        bank_id: str,
-        query: InterleavedTextMedia,
-        params: Optional[Dict[str, Any]] = None,
-    ) -> QueryDocumentsResponse:
-        index = self.cache.get(bank_id)
-        if index is None:
-            raise ValueError(f"Bank {bank_id} not found")
-
-        return await index.query_documents(query, params)

llama_stack/providers/impls/meta_reference/safety/__init__.py

@@ -1,17 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from .config import SafetyConfig
-
-
-async def get_provider_impl(config: SafetyConfig, deps):
-    from .safety import MetaReferenceSafetyImpl
-
-    assert isinstance(config, SafetyConfig), f"Unexpected config type: {type(config)}"
-
-    impl = MetaReferenceSafetyImpl(config, deps)
-    await impl.initialize()
-    return impl

llama_stack/providers/impls/meta_reference/safety/base.py

@@ -1,57 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from abc import ABC, abstractmethod
-from typing import List
-
-from llama_models.llama3.api.datatypes import interleaved_text_media_as_str, Message
-from pydantic import BaseModel
-from llama_stack.apis.safety import *  # noqa: F403
-
-CANNED_RESPONSE_TEXT = "I can't answer that. Can I help with something else?"
-
-
-# TODO: clean this up; just remove this type completely
-class ShieldResponse(BaseModel):
-    is_violation: bool
-    violation_type: Optional[str] = None
-    violation_return_message: Optional[str] = None
-
-
-# TODO: this is a caller / agent concern
-class OnViolationAction(Enum):
-    IGNORE = 0
-    WARN = 1
-    RAISE = 2
-
-
-class ShieldBase(ABC):
-    def __init__(
-        self,
-        on_violation_action: OnViolationAction = OnViolationAction.RAISE,
-    ):
-        self.on_violation_action = on_violation_action
-
-    @abstractmethod
-    async def run(self, messages: List[Message]) -> ShieldResponse:
-        raise NotImplementedError()
-
-
-def message_content_as_str(message: Message) -> str:
-    return interleaved_text_media_as_str(message.content)
-
-
-class TextShield(ShieldBase):
-    def convert_messages_to_text(self, messages: List[Message]) -> str:
-        return "\n".join([message_content_as_str(m) for m in messages])
-
-    async def run(self, messages: List[Message]) -> ShieldResponse:
-        text = self.convert_messages_to_text(messages)
-        return await self.run_impl(text)
-
-    @abstractmethod
-    async def run_impl(self, text: str) -> ShieldResponse:
-        raise NotImplementedError()

llama_stack/providers/impls/meta_reference/safety/config.py

@@ -1,48 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from enum import Enum
-from typing import List, Optional
-
-from llama_models.sku_list import CoreModelId, safety_models
-
-from pydantic import BaseModel, field_validator
-
-
-class PromptGuardType(Enum):
-    injection = "injection"
-    jailbreak = "jailbreak"
-
-
-class LlamaGuardShieldConfig(BaseModel):
-    model: str = "Llama-Guard-3-1B"
-    excluded_categories: List[str] = []
-
-    @field_validator("model")
-    @classmethod
-    def validate_model(cls, model: str) -> str:
-        permitted_models = [
-            m.descriptor()
-            for m in safety_models()
-            if (
-                m.core_model_id
-                in {
-                    CoreModelId.llama_guard_3_8b,
-                    CoreModelId.llama_guard_3_1b,
-                    CoreModelId.llama_guard_3_11b_vision,
-                }
-            )
-        ]
-        if model not in permitted_models:
-            raise ValueError(
-                f"Invalid model: {model}. Must be one of {permitted_models}"
-            )
-        return model
-
-
-class SafetyConfig(BaseModel):
-    llama_guard_shield: Optional[LlamaGuardShieldConfig] = None
-    enable_prompt_guard: Optional[bool] = False

llama_stack/providers/impls/meta_reference/safety/llama_guard.py

@@ -1,268 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-import re
-
-from string import Template
-from typing import List, Optional
-
-from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.inference import *  # noqa: F403
-
-from .base import CANNED_RESPONSE_TEXT, OnViolationAction, ShieldBase, ShieldResponse
-
-
-SAFE_RESPONSE = "safe"
-_INSTANCE = None
-
-CAT_VIOLENT_CRIMES = "Violent Crimes"
-CAT_NON_VIOLENT_CRIMES = "Non-Violent Crimes"
-CAT_SEX_CRIMES = "Sex Crimes"
-CAT_CHILD_EXPLOITATION = "Child Exploitation"
-CAT_DEFAMATION = "Defamation"
-CAT_SPECIALIZED_ADVICE = "Specialized Advice"
-CAT_PRIVACY = "Privacy"
-CAT_INTELLECTUAL_PROPERTY = "Intellectual Property"
-CAT_INDISCRIMINATE_WEAPONS = "Indiscriminate Weapons"
-CAT_HATE = "Hate"
-CAT_SELF_HARM = "Self-Harm"
-CAT_SEXUAL_CONTENT = "Sexual Content"
-CAT_ELECTIONS = "Elections"
-CAT_CODE_INTERPRETER_ABUSE = "Code Interpreter Abuse"
-
-
-SAFETY_CATEGORIES_TO_CODE_MAP = {
-    CAT_VIOLENT_CRIMES: "S1",
-    CAT_NON_VIOLENT_CRIMES: "S2",
-    CAT_SEX_CRIMES: "S3",
-    CAT_CHILD_EXPLOITATION: "S4",
-    CAT_DEFAMATION: "S5",
-    CAT_SPECIALIZED_ADVICE: "S6",
-    CAT_PRIVACY: "S7",
-    CAT_INTELLECTUAL_PROPERTY: "S8",
-    CAT_INDISCRIMINATE_WEAPONS: "S9",
-    CAT_HATE: "S10",
-    CAT_SELF_HARM: "S11",
-    CAT_SEXUAL_CONTENT: "S12",
-    CAT_ELECTIONS: "S13",
-    CAT_CODE_INTERPRETER_ABUSE: "S14",
-}
-
-
-DEFAULT_LG_V3_SAFETY_CATEGORIES = [
-    CAT_VIOLENT_CRIMES,
-    CAT_NON_VIOLENT_CRIMES,
-    CAT_SEX_CRIMES,
-    CAT_CHILD_EXPLOITATION,
-    CAT_DEFAMATION,
-    CAT_SPECIALIZED_ADVICE,
-    CAT_PRIVACY,
-    CAT_INTELLECTUAL_PROPERTY,
-    CAT_INDISCRIMINATE_WEAPONS,
-    CAT_HATE,
-    CAT_SELF_HARM,
-    CAT_SEXUAL_CONTENT,
-    CAT_ELECTIONS,
-]
-
-
-MODEL_TO_SAFETY_CATEGORIES_MAP = {
-    CoreModelId.llama_guard_3_8b.value: (
-        DEFAULT_LG_V3_SAFETY_CATEGORIES + [CAT_CODE_INTERPRETER_ABUSE]
-    ),
-    CoreModelId.llama_guard_3_1b.value: DEFAULT_LG_V3_SAFETY_CATEGORIES,
-    CoreModelId.llama_guard_3_11b_vision.value: DEFAULT_LG_V3_SAFETY_CATEGORIES,
-}
-
-
-PROMPT_TASK = "Task: Check if there is unsafe content in '$agent_type' messages in conversations according our safety policy with the below categories."
-
-SAFETY_CATEGORIES = """
-
-<BEGIN UNSAFE CONTENT CATEGORIES>
-$categories
-<END UNSAFE CONTENT CATEGORIES>
-"""
-
-PROMPT_CONVERSATION = """
-<BEGIN CONVERSATION>
-
-$conversations
-
-<END CONVERSATION>
-"""
-
-
-PROMPT_INSTRUCTIONS = """
-Provide your safety assessment for ONLY THE LAST $agent_type message in the above conversation:
- - First line must read 'safe' or 'unsafe'.
- - If unsafe, a second line must include a comma-separated list of violated categories."""
-
-
-PROMPT_TEMPLATE = Template(
-    f"{PROMPT_TASK}{SAFETY_CATEGORIES}{PROMPT_CONVERSATION}{PROMPT_INSTRUCTIONS}"
-)
-
-
-class LlamaGuardShield(ShieldBase):
-    def __init__(
-        self,
-        model: str,
-        inference_api: Inference,
-        excluded_categories: List[str] = None,
-        on_violation_action: OnViolationAction = OnViolationAction.RAISE,
-    ):
-        super().__init__(on_violation_action)
-
-        if excluded_categories is None:
-            excluded_categories = []
-
-        assert len(excluded_categories) == 0 or all(
-            x in SAFETY_CATEGORIES_TO_CODE_MAP.values() for x in excluded_categories
-        ), "Invalid categories in excluded categories. Expected format is ['S1', 'S2', ..]"
-
-        if model not in MODEL_TO_SAFETY_CATEGORIES_MAP:
-            raise ValueError(f"Unsupported model: {model}")
-
-        self.model = model
-        self.inference_api = inference_api
-        self.excluded_categories = excluded_categories
-
-    def check_unsafe_response(self, response: str) -> Optional[str]:
-        match = re.match(r"^unsafe\n(.*)$", response)
-        if match:
-            # extracts the unsafe code
-            extracted = match.group(1)
-            return extracted
-
-        return None
-
-    def get_safety_categories(self) -> List[str]:
-        excluded_categories = self.excluded_categories
-        if set(excluded_categories) == set(SAFETY_CATEGORIES_TO_CODE_MAP.values()):
-            excluded_categories = []
-
-        final_categories = []
-
-        all_categories = MODEL_TO_SAFETY_CATEGORIES_MAP[self.model]
-        for cat in all_categories:
-            cat_code = SAFETY_CATEGORIES_TO_CODE_MAP[cat]
-            if cat_code in excluded_categories:
-                continue
-            final_categories.append(f"{cat_code}: {cat}.")
-
-        return final_categories
-
-    def validate_messages(self, messages: List[Message]) -> None:
-        if len(messages) == 0:
-            raise ValueError("Messages must not be empty")
-        if messages[0].role != Role.user.value:
-            raise ValueError("Messages must start with user")
-
-        if len(messages) >= 2 and (
-            messages[0].role == Role.user.value and messages[1].role == Role.user.value
-        ):
-            messages = messages[1:]
-
-        for i in range(1, len(messages)):
-            if messages[i].role == messages[i - 1].role:
-                raise ValueError(
-                    f"Messages must alternate between user and assistant. Message {i} has the same role as message {i - 1}"
-                )
-        return messages
-
-    async def run(self, messages: List[Message]) -> ShieldResponse:
-        messages = self.validate_messages(messages)
-
-        if self.model == CoreModelId.llama_guard_3_11b_vision.value:
-            shield_input_message = self.build_vision_shield_input(messages)
-        else:
-            shield_input_message = self.build_text_shield_input(messages)
-
-        # TODO: llama-stack inference protocol has issues with non-streaming inference code
-        content = ""
-        async for chunk in self.inference_api.chat_completion(
-            model=self.model,
-            messages=[shield_input_message],
-            stream=True,
-        ):
-            event = chunk.event
-            if event.event_type == ChatCompletionResponseEventType.progress:
-                assert isinstance(event.delta, str)
-                content += event.delta
-
-        content = content.strip()
-        shield_response = self.get_shield_response(content)
-        return shield_response
-
-    def build_text_shield_input(self, messages: List[Message]) -> UserMessage:
-        return UserMessage(content=self.build_prompt(messages))
-
-    def build_vision_shield_input(self, messages: List[Message]) -> UserMessage:
-        conversation = []
-        most_recent_img = None
-
-        for m in messages[::-1]:
-            if isinstance(m.content, str):
-                conversation.append(m)
-            elif isinstance(m.content, ImageMedia):
-                if most_recent_img is None and m.role == Role.user.value:
-                    most_recent_img = m.content
-                    conversation.append(m)
-            elif isinstance(m.content, list):
-                content = []
-                for c in m.content:
-                    if isinstance(c, str):
-                        content.append(c)
-                    elif isinstance(c, ImageMedia):
-                        if most_recent_img is None and m.role == Role.user.value:
-                            most_recent_img = c
-                            content.append(c)
-                    else:
-                        raise ValueError(f"Unknown content type: {c}")
-
-                conversation.append(UserMessage(content=content))
-            else:
-                raise ValueError(f"Unknown content type: {m.content}")
-
-        prompt = []
-        if most_recent_img is not None:
-            prompt.append(most_recent_img)
-        prompt.append(self.build_prompt(conversation[::-1]))
-
-        return UserMessage(content=prompt)
-
-    def build_prompt(self, messages: List[Message]) -> str:
-        categories = self.get_safety_categories()
-        categories_str = "\n".join(categories)
-        conversations_str = "\n\n".join(
-            [
-                f"{m.role.capitalize()}: {interleaved_text_media_as_str(m.content)}"
-                for m in messages
-            ]
-        )
-        return PROMPT_TEMPLATE.substitute(
-            agent_type=messages[-1].role.capitalize(),
-            categories=categories_str,
-            conversations=conversations_str,
-        )
-
-    def get_shield_response(self, response: str) -> ShieldResponse:
-        response = response.strip()
-        if response == SAFE_RESPONSE:
-            return ShieldResponse(is_violation=False)
-        unsafe_code = self.check_unsafe_response(response)
-        if unsafe_code:
-            unsafe_code_list = unsafe_code.split(",")
-            if set(unsafe_code_list).issubset(set(self.excluded_categories)):
-                return ShieldResponse(is_violation=False)
-            return ShieldResponse(
-                is_violation=True,
-                violation_type=unsafe_code,
-                violation_return_message=CANNED_RESPONSE_TEXT,
-            )
-
-        raise ValueError(f"Unexpected response: {response}")

llama_stack/providers/impls/meta_reference/safety/prompt_guard.py

@@ -1,145 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the terms described in the LICENSE file in
-# the root directory of this source tree.
-
-from enum import auto, Enum
-from typing import List
-
-import torch
-
-from llama_models.llama3.api.datatypes import Message
-from termcolor import cprint
-
-from .base import message_content_as_str, OnViolationAction, ShieldResponse, TextShield
-
-
-class PromptGuardShield(TextShield):
-    class Mode(Enum):
-        INJECTION = auto()
-        JAILBREAK = auto()
-
-    _instances = {}
-    _model_cache = None
-
-    @staticmethod
-    def instance(
-        model_dir: str,
-        threshold: float = 0.9,
-        temperature: float = 1.0,
-        mode: "PromptGuardShield.Mode" = Mode.JAILBREAK,
-        on_violation_action=OnViolationAction.RAISE,
-    ) -> "PromptGuardShield":
-        action_value = on_violation_action.value
-        key = (model_dir, threshold, temperature, mode, action_value)
-        if key not in PromptGuardShield._instances:
-            PromptGuardShield._instances[key] = PromptGuardShield(
-                model_dir=model_dir,
-                threshold=threshold,
-                temperature=temperature,
-                mode=mode,
-                on_violation_action=on_violation_action,
-            )
-        return PromptGuardShield._instances[key]
-
-    def __init__(
-        self,
-        model_dir: str,
-        threshold: float = 0.9,
-        temperature: float = 1.0,
-        mode: "PromptGuardShield.Mode" = Mode.JAILBREAK,
-        on_violation_action: OnViolationAction = OnViolationAction.RAISE,
-    ):
-        super().__init__(on_violation_action)
-        assert (
-            model_dir is not None
-        ), "Must provide a model directory for prompt injection shield"
-        if temperature <= 0:
-            raise ValueError("Temperature must be greater than 0")
-        self.device = "cuda"
-        if PromptGuardShield._model_cache is None:
-            from transformers import AutoModelForSequenceClassification, AutoTokenizer
-
-            # load model and tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(model_dir)
-            model = AutoModelForSequenceClassification.from_pretrained(
-                model_dir, device_map=self.device
-            )
-            PromptGuardShield._model_cache = (tokenizer, model)
-
-        self.tokenizer, self.model = PromptGuardShield._model_cache
-        self.temperature = temperature
-        self.threshold = threshold
-        self.mode = mode
-
-    def convert_messages_to_text(self, messages: List[Message]) -> str:
-        return message_content_as_str(messages[-1])
-
-    async def run_impl(self, text: str) -> ShieldResponse:
-        # run model on messages and return response
-        inputs = self.tokenizer(text, return_tensors="pt")
-        inputs = {name: tensor.to(self.model.device) for name, tensor in inputs.items()}
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        logits = outputs[0]
-        probabilities = torch.softmax(logits / self.temperature, dim=-1)
-        score_embedded = probabilities[0, 1].item()
-        score_malicious = probabilities[0, 2].item()
-        cprint(
-            f"Ran PromptGuardShield and got Scores: Embedded: {score_embedded}, Malicious: {score_malicious}",
-            color="magenta",
-        )
-
-        if self.mode == self.Mode.INJECTION and (
-            score_embedded + score_malicious > self.threshold
-        ):
-            return ShieldResponse(
-                is_violation=True,
-                violation_type=f"prompt_injection:embedded={score_embedded},malicious={score_malicious}",
-                violation_return_message="Sorry, I cannot do this.",
-            )
-        elif self.mode == self.Mode.JAILBREAK and score_malicious > self.threshold:
-            return ShieldResponse(
-                is_violation=True,
-                violation_type=f"prompt_injection:malicious={score_malicious}",
-                violation_return_message="Sorry, I cannot do this.",
-            )
-
-        return ShieldResponse(
-            is_violation=False,
-        )
-
-
-class JailbreakShield(PromptGuardShield):
-    def __init__(
-        self,
-        model_dir: str,
-        threshold: float = 0.9,
-        temperature: float = 1.0,
-        on_violation_action: OnViolationAction = OnViolationAction.RAISE,
-    ):
-        super().__init__(
-            model_dir=model_dir,
-            threshold=threshold,
-            temperature=temperature,
-            mode=PromptGuardShield.Mode.JAILBREAK,
-            on_violation_action=on_violation_action,
-        )
-
-
-class InjectionShield(PromptGuardShield):
-    def __init__(
-        self,
-        model_dir: str,
-        threshold: float = 0.9,
-        temperature: float = 1.0,
-        on_violation_action: OnViolationAction = OnViolationAction.RAISE,
-    ):
-        super().__init__(
-            model_dir=model_dir,
-            threshold=threshold,
-            temperature=temperature,
-            mode=PromptGuardShield.Mode.INJECTION,
-            on_violation_action=on_violation_action,
-        )