diff --git a/docs/devel_doc/openapi.json b/docs/devel_doc/openapi.json index eed3b377f..33c1f655a 100644 --- a/docs/devel_doc/openapi.json +++ b/docs/devel_doc/openapi.json @@ -14639,7 +14639,13 @@ "url": { "type": "string", "title": "Base URL", - "description": "The model_id to use for the guard" + "description": "Base URL of the OpenAI-compatible inference endpoint." + }, + "model_id": { + "type": "string", + "title": "Model name", + "description": "Model name sent to the inference server. Override when the server registers the model under a different name (e.g. an Ollama tag). The prompt template is built for the 4.1 format.", + "default": "ibm-granite/granite-guardian-4.1-8b" }, "api_key": { "anyOf": [ @@ -14686,6 +14692,21 @@ "description": "SSL certificate verification. Can be:\n - True: Verify using system CA bundle (default, recommended)\n - False: Disable verification (insecure, for dev only)\n - str: Path to custom CA bundle file (for internal PKI)", "default": true }, + "parallel": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "integer", + "maximum": 10.0, + "minimum": 1.0 + } + ], + "title": "Parallel execution", + "description": "True to run all risk checks in parallel, False to run sequentially, or an integer 1-10 for explicit batch size.", + "default": 3 + }, "risks": { "items": { "$ref": "#/components/schemas/RiskDefinition" @@ -14733,6 +14754,14 @@ "title": "GraniteGuardianShieldConfiguration", "description": "Configuration for a named Granite Guardian guardrail shield.\n\nAttributes:\n name: Unique, user-facing name identifying this shield instance.\n provider_id: Discriminator identifying this as a granite-guardian shield.\n config: Granite-guardian-specific configuration." }, + "GuardrailPoint": { + "type": "string", + "enum": [ + "input", + "output", + "tool" + ] + }, "HTTPAuthSecurityScheme": { "properties": { "bearerFormat": { @@ -20914,12 +20943,7 @@ }, "points": { "items": { - "type": "string", - "enum": [ - "input", - "output", - "tool" - ] + "$ref": "#/components/schemas/GuardrailPoint" }, "type": "array", "minItems": 1, diff --git a/src/models/config.py b/src/models/config.py index 66825ed86..4c393c62f 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -23,6 +23,7 @@ PositiveInt, PrivateAttr, SecretStr, + StrictBool, field_validator, model_validator, ) @@ -34,6 +35,8 @@ from utils.mcp_auth_headers import resolve_authorization_headers from utils.types import CompiledPatterns +type GuardrailPoint = Literal["input", "output", "tool"] + logger = get_logger(__name__) @@ -2422,8 +2425,7 @@ def validate_providers_and_default(self) -> Self: if self.default_provider is None: raise ValueError( - "vector_store.default_provider is required when providers " - "is non-empty" + "vector_store.default_provider is required when providers is non-empty" ) ids = [provider.id for provider in self.providers] @@ -3197,7 +3199,7 @@ class RiskDefinition(ConfigurationBase): "reasoning before scoring." ), ) - points: list[Literal["input", "output", "tool"]] = Field( + points: list[GuardrailPoint] = Field( ..., min_length=1, title="Guardrail points", @@ -3217,7 +3219,19 @@ class GraniteGuardianConfig(ConfigurationBase): """Configuration for the Granite Guardian moderation guardrail.""" url: str = Field( - ..., title="Base URL", description="The model_id to use for the guard" + ..., + title="Base URL", + description="Base URL of the OpenAI-compatible inference endpoint.", + ) + + model_id: str = Field( + "ibm-granite/granite-guardian-4.1-8b", + title="Model name", + description=( + "Model name sent to the inference server. Override when the " + "server registers the model under a different name (e.g. an " + "Ollama tag). The prompt template is built for the 4.1 format." + ), ) api_key: Optional[SecretStr] = Field( @@ -3243,6 +3257,16 @@ class GraniteGuardianConfig(ConfigurationBase): ), ) + parallel: StrictBool | Annotated[int, Field(ge=1, le=10)] = Field( + default=3, + title="Parallel execution", + description=( + "True to run all risk checks in parallel, " + "False to run sequentially, " + "or an integer 1-10 for explicit batch size." + ), + ) + risks: list[RiskDefinition] = Field( ..., title="Defined risks", diff --git a/src/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py new file mode 100644 index 000000000..48e7db6b7 --- /dev/null +++ b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py @@ -0,0 +1,7 @@ +"""Granite Guardian safety capability for risk-based moderation.""" + +from pydantic_ai_lightspeed.capabilities.granite_guardian._capability import ( + GraniteGuardian, +) + +__all__ = ["GraniteGuardian"] diff --git a/src/pydantic_ai_lightspeed/capabilities/granite_guardian/_capability.py b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/_capability.py new file mode 100644 index 000000000..e73c409fe --- /dev/null +++ b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/_capability.py @@ -0,0 +1,313 @@ +"""Granite Guardian safety capability for input/output guardrail moderation.""" + +import asyncio +from dataclasses import dataclass, field +from typing import ClassVar, Optional +from uuid import uuid4 + +import httpx +from openai import AsyncOpenAI +from pydantic import StrictBool +from pydantic_ai import AgentRunResult, RunContext +from pydantic_ai._agent_graph import GraphAgentState +from pydantic_ai.capabilities import WrapRunHandler +from pydantic_ai.direct import model_request +from pydantic_ai.exceptions import UnexpectedModelBehavior +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, +) +from pydantic_ai.models import Model +from pydantic_ai.models.openai import OpenAIChatModel, OpenAIChatModelSettings +from pydantic_ai.providers.openai import OpenAIProvider +from pydantic_ai.usage import RequestUsage + +from client.ogx import AsyncOgxClientHolder +from log import get_logger +from models.common.moderation import ( + ShieldModerationBlocked, + ShieldModerationPassed, + ShieldModerationResult, +) +from models.config import GraniteGuardianConfig, GuardrailPoint, RiskDefinition +from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability +from pydantic_ai_lightspeed.capabilities.granite_guardian.utils import ( + build_guardian_block, + is_safe, +) +from pydantic_ai_lightspeed.capabilities.utils import ( + extract_conversation_id, + message_to_str, +) +from utils.conversations import append_turn_to_conversation + +type Guardrail = tuple[str, str, float, str] + +logger = get_logger(__name__) + + +async def _package_risk_check_task( + prompt: str, guardrail: Guardrail, model: Model +) -> tuple[ModelResponse, float, str]: + """Run a single Guardian risk check and return its result. + + Parameters: + prompt: The text to evaluate. + guardrail: A guardrail tuple of (name, block, threshold, violation_message). + model: The Granite Guardian model to use for evaluation. + + Returns: + A tuple of (model_response, threshold, violation_message). + """ + name, block, threshold, violation_message = guardrail + start = asyncio.get_event_loop().time() + result = await model_request( + model=model, + messages=[ModelRequest.user_text_prompt(prompt, instructions=block)], + model_settings=OpenAIChatModelSettings( + openai_logprobs=True, openai_top_logprobs=20 + ), + ) + elapsed = asyncio.get_event_loop().time() - start + logger.info("Guardian risk '%s' completed in %.3fs", name, elapsed) + return result, threshold, violation_message + + +async def _run_risk_check( + prompt: str, + model: Model, + guardrails: list[Guardrail], + batch_size: int = 3, +) -> tuple[Optional[str], RequestUsage]: + """Evaluate the prompt against guardrails in parallel batches. + + Guardrails are dispatched concurrently in batches. Within each batch, + all checks run in parallel; if any violation is found the remaining + batches are skipped. Per-rule latency is logged at INFO level. + + Parameters: + prompt: The text to evaluate. + model: The Granite Guardian model to use for evaluation. + guardrails: Ordered list of guardrail tuples to check. + batch_size: Number of risk checks to run in parallel per batch. + + Returns: + A tuple of (violation_message, token_usage). violation_message is + None when all checks pass. + + Raises: + UnexpectedModelBehavior: When the model response is missing + provider_details or logprobs. + """ + token_usage = RequestUsage() + + for i in range(0, len(guardrails), batch_size): + batch = [ + _package_risk_check_task(prompt, g, model) + for g in guardrails[i : i + batch_size] + ] + + results = await asyncio.gather(*batch) + + for result, _, _ in results: + token_usage.incr(result.usage) + + for result, threshold, violation_message in results: + if not result.provider_details: + raise UnexpectedModelBehavior( + "No provider_details provided from granite guardian's response" + ) + + logprobs = result.provider_details.get("logprobs") + if not logprobs: + raise UnexpectedModelBehavior("No logprobs field in provider_details") + + if not is_safe(threshold, logprobs): + return violation_message, token_usage + + return None, token_usage + + +def _filter_guardrails( + risks: list[RiskDefinition], point: GuardrailPoint +) -> list[Guardrail]: + """Filter risk definitions to guardrail tuples for a given guardrail point. + + Parameters: + risks: All configured risk definitions. + point: The guardrail point to filter by (INPUT, OUTPUT, or TOOL). + + Returns: + A list of guardrail tuples for enabled risks matching the point. + """ + return [ + ( + risk.name, + build_guardian_block(risk.description, think=risk.enable_thinking), + risk.threshold, + risk.violation_message, + ) + for risk in risks + if risk.enabled and point in risk.points + ] + + +def _get_batch_size(parallel: StrictBool | int, num_guardrail: int) -> int: + """Resolve the parallel setting to a concrete batch size. + + Parameters: + parallel: True for full parallelism, False for sequential, or an + explicit batch size. + num_guardrail: Total number of guardrails to run. + + Returns: + The number of risk checks to run concurrently per batch. + """ + if isinstance(parallel, bool): + return max(1, num_guardrail) if parallel else 1 + + return parallel + + +@dataclass +class GraniteGuardian(AbstractSafetyCapability): + """Safety capability using Granite Guardian for risk-based moderation. + + Uses Granite Guardian's logprob-based scoring to evaluate user input + against configured risk categories. When used as a pydantic-ai capability, + ``wrap_run`` applies input guardrails. The ``run`` method provides a + standalone shield interface for use outside the agent lifecycle. + + Attributes: + config: Granite Guardian configuration with risks and connection details. + run_moderation_guardrail_point: The guardrail point used by the + standalone ``run`` method. + """ + + config: GraniteGuardianConfig + run_moderation_guardrail_point: GuardrailPoint = "input" + _model: Model = field(init=False) + # Only one Granite Guardian shield should be configured; multiple entries are + # unsupported. A dict is used defensively so that a misconfiguration with two + # distinct configs does not cause one to silently overwrite the other's model. + _model_cache: ClassVar[dict[int, Model]] = {} + + def __post_init__(self) -> None: + """Initialize the Granite Guardian model with the configured provider.""" + cache_key = id(self.config) + if cache_key in GraniteGuardian._model_cache: + self._model = GraniteGuardian._model_cache[cache_key] + return + + http_client = httpx.AsyncClient( + verify=self.config.verify_ssl, + timeout=self.config.timeout, + ) + + # When we attach the API key to the request, we need to make sure we encrypt the + # request by communicating through https + base_url = httpx.URL(self.config.url) + if self.config.api_key is not None and base_url.scheme != "https": + raise ValueError( + "Granite Guardian endpoints with an API key must use HTTPS" + ) + + openai_client = AsyncOpenAI( + base_url=self.config.url, + api_key=( + self.config.api_key.get_secret_value() # pylint: disable=no-member + if self.config.api_key is not None + else "api-key-not-set" + ), + max_retries=self.config.max_retries, + http_client=http_client, + ) + + provider = OpenAIProvider(openai_client=openai_client) + + self._model = OpenAIChatModel(self.config.model_id, provider=provider) + GraniteGuardian._model_cache[cache_key] = self._model + + async def wrap_run( + self, ctx: RunContext, *, handler: WrapRunHandler + ) -> AgentRunResult: + """Apply input guardrails before the agent run. + + Evaluates the user prompt against all INPUT-point risks. If any risk + is violated, the run is short-circuited with a rejection message. + Otherwise, the handler is called to proceed with the real run. + + Parameters: + ctx: The run context containing the user prompt and usage tracker. + handler: The handler to call if the input passes all guardrails. + + Returns: + The agent run result, either a rejection or the handler's result. + """ + user_prompt = message_to_str(ctx.prompt) + + input_guardrails = _filter_guardrails(self.config.risks, "input") + batch_size = _get_batch_size(self.config.parallel, len(input_guardrails)) + # TODO: We need to consider how we want to reveal the token usage for Granite Guardian, # pylint: disable=fixme + # since combining the token usage with the main inference model is not a right thing to do. + violation_message, _ = await _run_risk_check( + user_prompt, self._model, input_guardrails, batch_size + ) + + if violation_message is not None: + state = GraphAgentState( + usage=ctx.usage, + message_history=[ + ModelRequest.user_text_prompt(user_prompt), + ModelResponse( + [TextPart(violation_message)], + finish_reason="stop", + ), + ], + ) + + conversation_id = extract_conversation_id(ctx.model) + if conversation_id is not None: + await append_turn_to_conversation( + AsyncOgxClientHolder().get_client(), + conversation_id, + user_prompt, + violation_message, + ) + else: + logger.warning( + "Unable to determine conversation ID from model settings; " + "skipping v1/conversation persistence for rejected question." + ) + + return AgentRunResult(output=violation_message, _state=state) + + return await handler() # proceed with the real run + + async def run(self, input_text: str) -> ShieldModerationResult: + """Run standalone shield moderation on the given text. + + Uses ``run_moderation_guardrail_point`` to filter which risks apply. + + Parameters: + input_text: The text to evaluate. + + Returns: + A blocked result with the violation message, or a passed result. + """ + filtered_guardrails = _filter_guardrails( + self.config.risks, self.run_moderation_guardrail_point + ) + batch_size = _get_batch_size(self.config.parallel, len(filtered_guardrails)) + + violation_message, _ = await _run_risk_check( + input_text, self._model, filtered_guardrails, batch_size + ) + + if violation_message is not None: + return ShieldModerationBlocked( + message=violation_message, moderation_id=f"modr-{uuid4()}" + ) + + return ShieldModerationPassed() diff --git a/src/pydantic_ai_lightspeed/capabilities/granite_guardian/utils.py b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/utils.py new file mode 100644 index 000000000..44b75d15f --- /dev/null +++ b/src/pydantic_ai_lightspeed/capabilities/granite_guardian/utils.py @@ -0,0 +1,222 @@ +"""Utility functions for Granite Guardian prompt construction and logprob parsing.""" + +from enum import Enum, StrEnum, auto +from math import exp +from typing import Any + +from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob +from pydantic_ai.exceptions import UnexpectedModelBehavior + +from log import get_logger + + +class ParsingStatus(Enum): + """State machine states for parsing Granite Guardian response tokens.""" + + BEFORE_THINK = auto() + INSIDE_THINK = auto() + POST_THINK = auto() + INSIDE_SCORE = auto() + + +class Token(StrEnum): + """Granite Guardian verdict tokens indicating safe or risky content.""" + + SAFE = "no" + RISKY = "yes" + + +THINK_START_TAG = "" +THINK_END_TAG = "" +SCORE_START_TAG = "" +SCORE_END_TAG = "" + +# Copied the implementation of Granite Guardian guardrail from Ask Red Hat. +# Reference: https://gitlab.cee.redhat.com/ask-red-hat/backend/ask-guardrails + +# ============================================================================= +# Guardian Block Construction (GG 4.1 client-side prompt) +# ============================================================================= + +_JUDGE_PREAMBLE_NOTHINK = ( + "As a judge agent, assess whether the provided text " + "meets the given judging criteria using all available information, " + "including conversations, documents, and tools. Provide your score " + "immediately without explanation. Output empty \\n tags " + "followed by your score in tags." +) + +_JUDGE_PREAMBLE_THINK = ( + "As a judge agent, carefully analyze whether the provided " + "text meets the given judging criteria. Use all available information, " + "including conversations, documents, and tools. Reason through the " + "evidence step by step before providing your score. Wrap your reasoning " + "in tags, then provide your score in tags." +) + +_SCORING_SCHEMA_PROMPT_ONLY = ( + "If the last user's text meets the criteria, return 'yes'; otherwise, return 'no'." +) + +logger = get_logger(__name__) + + +def build_guardian_block(criteria: str, *, think: bool = False) -> str: + """Build the guardian evaluation block sent as a second user message.""" + preamble = _JUDGE_PREAMBLE_THINK if think else _JUDGE_PREAMBLE_NOTHINK + return ( + f"{preamble}\n\n" + f"### Criteria: {criteria}\n\n" + f"### Scoring Schema: {_SCORING_SCHEMA_PROMPT_ONLY}" + ) + + +def _search_tag(tag: str, buffer: str) -> tuple[str, bool]: + """Search for a complete XML tag in the buffer. + + Parameters: + tag: The XML tag to search for (e.g. ````). + buffer: The accumulated token text to search in. + + Returns: + A tuple of (remaining_buffer, found). The buffer is cleared when + the content is not tag-like or when the tag is found. + """ + if not buffer.strip().startswith("<"): + return "", False + + if tag in buffer: + return "", True + + return buffer, False + + +def _clean_up_candidates(candidates: list[ChatCompletionTokenLogprob]) -> None: + """Remove trailing candidates that are part of the score end tag. + + Pops tokens from the end of the candidates list until the ```` + tag boundary is found, leaving only the score content tokens. + + Parameters: + candidates: Mutable list of token logprobs to clean in place. + """ + buffer = "" + while len(candidates) != 0: + cur_candidate = candidates.pop() + buffer = cur_candidate.token + buffer + if SCORE_END_TAG in buffer: + return + + +def _extract_tokens_inside_score_tag( + logprobs: list[dict[str, Any]], +) -> ChatCompletionTokenLogprob: + """Extract the single token logprob from inside the ```` tag. + + Parses the model's structured output through ```` and ```` + tags, returning the logprob of the scoring token (``yes`` or ``no``). + + Parameters: + logprobs: Raw logprob dictionaries from the model's provider_details. + + Returns: + The ``ChatCompletionTokenLogprob`` for the score token. + + Raises: + UnexpectedModelBehavior: When the model output doesn't contain the + expected tag structure, or the score tag has zero or multiple tokens. + """ + cur_buffer = "" + cur_status = ParsingStatus.BEFORE_THINK + candidates: list[ChatCompletionTokenLogprob] = [] + + for _logprob in logprobs: + logprob = ChatCompletionTokenLogprob.model_validate(_logprob) + cur_buffer += logprob.token + + match cur_status: + case ParsingStatus.BEFORE_THINK: + # Sometimes Granite Guardian model will start with the end thinking tag + # if thinking is disabled + cur_buffer, found_end = _search_tag(THINK_END_TAG, cur_buffer) + if found_end: + cur_status = ParsingStatus.POST_THINK + continue + + cur_buffer, found_tag = _search_tag(THINK_START_TAG, cur_buffer) + if found_tag: + cur_status = ParsingStatus.INSIDE_THINK + case ParsingStatus.INSIDE_THINK: + cur_buffer, found_tag = _search_tag(THINK_END_TAG, cur_buffer) + if found_tag: + cur_status = ParsingStatus.POST_THINK + case ParsingStatus.POST_THINK: + cur_buffer, found_tag = _search_tag(SCORE_START_TAG, cur_buffer) + if found_tag: + cur_status = ParsingStatus.INSIDE_SCORE + case ParsingStatus.INSIDE_SCORE: + cur_buffer, found_tag = _search_tag(SCORE_END_TAG, cur_buffer) + candidates.append(logprob) + if found_tag: + _clean_up_candidates(candidates) + + # Theoretically, there should be only one token inside score tag because + # the chance that 'yes' and 'no' are split into two tokens are extremely + # low. Raising an error here so if it ever happens in the future, we can + # improve the parsing logic then. + candidates_num = len(candidates) + if candidates_num == 0: + raise UnexpectedModelBehavior("No token found inside score tag") + if candidates_num > 1: + raise UnexpectedModelBehavior( + "More than one token found inside score tag" + ) + return candidates[0] + + raise UnexpectedModelBehavior("Model did not generate the required format") + + +def _get_risky_probabilities(token: ChatCompletionTokenLogprob) -> float: + """Compute the normalized probability of the risky (``yes``) outcome. + + Parameters: + token: The score token with its top logprob alternatives. + + Returns: + The probability of the risky outcome, normalized against the sum + of safe and risky probabilities. + + Raises: + UnexpectedModelBehavior: When neither safe nor risky tokens appear + in the top logprobs (underflow). + """ + safe_prob, risky_prob = 0.0, 0.0 + for candidate_token in token.top_logprobs: + match candidate_token.token.strip().lower(): + case Token.SAFE: + safe_prob += exp(candidate_token.logprob) + case Token.RISKY: + risky_prob += exp(candidate_token.logprob) + total_prob = safe_prob + risky_prob + + if total_prob == 0: + raise UnexpectedModelBehavior("Logprob underflow") + + return risky_prob / total_prob + + +def is_safe(threshold: float, logprobs: list[dict[str, Any]]) -> bool: + """Determine whether the input is safe based on the risky probability. + + Parameters: + threshold: The risk threshold; the input is unsafe when the risky + probability meets or exceeds this value. + logprobs: Raw logprob dictionaries from the model response. + + Returns: + True if the risky probability is below the threshold (safe). + """ + token_for_confidence_threshold = _extract_tokens_inside_score_tag(logprobs) + p_risky = _get_risky_probabilities(token_for_confidence_threshold) + + return p_risky < threshold diff --git a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py index e26523aa1..2f6a0f754 100644 --- a/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py +++ b/src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py @@ -22,7 +22,6 @@ from pydantic_ai.messages import ( ModelRequest, ModelResponse, - TextContent, TextPart, UserContent, ) @@ -40,6 +39,10 @@ QuestionValidityConfig, ) from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability +from pydantic_ai_lightspeed.capabilities.utils import ( + extract_conversation_id, + message_to_str, +) from pydantic_ai_lightspeed.ogx import OgxResponsesModel from utils.conversations import append_turn_to_conversation @@ -49,67 +52,6 @@ SUBJECT_ALLOWED = "ALLOWED" -def _extract_message_str_from_user_content(user_content: Sequence[UserContent]) -> str: - """Extract and combine all text content into a string from a UserContent sequence. - - Parameters: - user_content: A sequence of user content items to extract text from. - - Returns: - A single string with all text content joined by newlines. - """ - str_arr: list[str] = [] - for c in user_content: - match c: - case str() as s: - str_arr.append(s) - case TextContent(content=c): - str_arr.append(c) - - return "\n".join(str_arr) - - -def _message_to_str(message: Optional[str | Sequence[UserContent]]) -> str: - """Convert a user message (string, content sequence, or None) to plain text. - - Parameters: - message: The user input as a string, sequence of user content, or None. - - Returns: - A plain-text representation of the message, or an empty string for None. - """ - match message: - case str() as s: - return s - case Sequence() as seq: - return _extract_message_str_from_user_content(seq) - case None: - return "" - - -def _extract_conversation_id(model: Model) -> Optional[str]: - """Extract the OGX conversation ID from the agent's model settings. - - The main agent's model is built with ``conversation`` in its - ``extra_body`` model settings (see ``OgxResponsesModel.from_ogx_client``). - This pulls it back out so the capability can persist the rejected turn - to the same conversation. - - Parameters: - model: The model bound to the current agent run (``ctx.model``). - - Returns: - The conversation ID, or None if the model has no such setting - (e.g. when used outside an OGX-backed agent). - """ - extra_body = (model.settings or {}).get("extra_body") - if not isinstance(extra_body, dict): - return None - - conversation_id = extra_body.get("conversation") - return conversation_id if isinstance(conversation_id, str) else None - - @dataclass class QuestionValidity(AbstractSafetyCapability): """Block or modify user input based on a guardrail check. @@ -149,7 +91,7 @@ def _build_prompt(self, message: Optional[str | Sequence[UserContent]]) -> str: The rendered prompt string ready to send to the validity model. """ return Template(self.config.model_prompt).substitute( - message=_message_to_str(message), + message=message_to_str(message), allowed=SUBJECT_ALLOWED, rejected=SUBJECT_REJECTED, ) @@ -185,7 +127,7 @@ async def wrap_run( return await handler() # proceed with the real run # short-circuit: return the rejection message with shield usage tracked - user_message = _message_to_str(ctx.prompt) + user_message = message_to_str(ctx.prompt) state = GraphAgentState( usage=ctx.usage, message_history=[ @@ -197,7 +139,7 @@ async def wrap_run( ], ) - conversation_id = _extract_conversation_id(ctx.model) + conversation_id = extract_conversation_id(ctx.model) if conversation_id is not None: await append_turn_to_conversation( AsyncOgxClientHolder().get_client(), diff --git a/src/pydantic_ai_lightspeed/capabilities/utils.py b/src/pydantic_ai_lightspeed/capabilities/utils.py new file mode 100644 index 000000000..2461fc6fa --- /dev/null +++ b/src/pydantic_ai_lightspeed/capabilities/utils.py @@ -0,0 +1,69 @@ +"""Shared utility functions for safety capabilities.""" + +from collections.abc import Sequence +from typing import Optional + +from pydantic_ai.messages import TextContent, UserContent +from pydantic_ai.models import Model + + +def extract_message_str_from_user_content( + user_content: Sequence[UserContent], +) -> str: + """Extract and combine all text content into a string from a UserContent sequence. + + Parameters: + user_content: A sequence of user content items to extract text from. + + Returns: + A single string with all text content joined by newlines. + """ + str_arr: list[str] = [] + for c in user_content: + match c: + case str() as s: + str_arr.append(s) + case TextContent(content=c): + str_arr.append(c) + + return "\n".join(str_arr) + + +def message_to_str(message: Optional[str | Sequence[UserContent]]) -> str: + """Convert a user message (string, content sequence, or None) to plain text. + + Parameters: + message: The user input as a string, sequence of user content, or None. + + Returns: + A plain-text representation of the message, or an empty string for None. + """ + match message: + case str() as s: + return s + case Sequence() as seq: + return extract_message_str_from_user_content(seq) + case None: + return "" + + +def extract_conversation_id(model: Model) -> Optional[str]: + """Extract the conversation ID from the agent's model settings. + + The main agent's model is built with ``conversation`` in its + ``extra_body`` model settings (see ``OgxResponsesModel.from_ogx_client``). + This pulls it back out so the capability can persist the rejected turn + to the same conversation. + + Parameters: + model: The model bound to the current agent run (``ctx.model``). + + Returns: + The conversation ID, or None if the model has no such setting. + """ + extra_body = (model.settings or {}).get("extra_body") + if not isinstance(extra_body, dict): + return None + + conversation_id = extra_body.get("conversation") + return conversation_id if isinstance(conversation_id, str) else None diff --git a/src/utils/pydantic_ai_helpers.py b/src/utils/pydantic_ai_helpers.py index d7ff1fa83..57861cc99 100644 --- a/src/utils/pydantic_ai_helpers.py +++ b/src/utils/pydantic_ai_helpers.py @@ -24,6 +24,9 @@ SkillsConfiguration, ) from pydantic_ai_lightspeed.capabilities import QuestionValidity +from pydantic_ai_lightspeed.capabilities.granite_guardian import ( + GraniteGuardian, +) from pydantic_ai_lightspeed.capabilities.redaction import PiiRedactionCapability from pydantic_ai_lightspeed.ogx import OgxResponsesModel from utils.shields import get_shields_for_request @@ -177,7 +180,7 @@ def _shield_capability(shield: ShieldConfiguration) -> AgentCapability[object]: case RedactionConfig(): return PiiRedactionCapability(config=shield.config) case GraniteGuardianConfig(): - raise NotImplementedError("Granite Guardian capability not implemented") + return GraniteGuardian(config=shield.config) case _: raise ValueError( f"Unsupported shield config type for shield '{shield.name}': " diff --git a/src/utils/shields.py b/src/utils/shields.py index 662e96304..80783a08d 100644 --- a/src/utils/shields.py +++ b/src/utils/shields.py @@ -23,11 +23,15 @@ ) from models.config import ( GraniteGuardianConfig, + GuardrailPoint, QuestionValidityConfig, RedactionConfig, ShieldConfiguration, ) from pydantic_ai_lightspeed.capabilities.base import AbstractSafetyCapability +from pydantic_ai_lightspeed.capabilities.granite_guardian import ( + GraniteGuardian, +) from pydantic_ai_lightspeed.capabilities.question_validity._capability import ( QuestionValidity, ) @@ -82,6 +86,7 @@ async def run_shield_moderation_v2( input_text: str, shield_configs: list[ShieldConfiguration], selected_shield_ids: Optional[list[str]] = None, + guardrail_point: GuardrailPoint = "input", ) -> ShieldModerationResult: """Run v2 shield moderation on input text. @@ -120,7 +125,7 @@ async def run_shield_moderation_v2( ) for shield_config in selected_shield_configs: - shield = build_shield(shield_config) + shield = build_shield(shield_config, guardrail_point) try: shield_result = await shield.run(input_text) @@ -149,7 +154,10 @@ async def run_shield_moderation_v2( return ShieldModerationPassed() -def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability: +def build_shield( + shield_config: ShieldConfiguration, + guardrail_point: GuardrailPoint = "input", +) -> AbstractSafetyCapability: """Build a safety capability instance from a shield configuration. Parameters: @@ -164,7 +172,7 @@ def build_shield(shield_config: ShieldConfiguration) -> AbstractSafetyCapability case RedactionConfig(): return PiiRedactionCapability(shield_config.config) case GraniteGuardianConfig(): - raise NotImplementedError("Granite Guardian capability not implemented") + return GraniteGuardian(shield_config.config, guardrail_point) case _: raise ValueError( f"Unsupported shield config type for shield '{shield_config.name}': " diff --git a/tests/integration/endpoints/test_granite_guardian_integration.py b/tests/integration/endpoints/test_granite_guardian_integration.py new file mode 100644 index 000000000..343595d2a --- /dev/null +++ b/tests/integration/endpoints/test_granite_guardian_integration.py @@ -0,0 +1,820 @@ +"""Integration tests for Granite Guardian input guardrail across endpoints. + +Tests exercise the shield moderation pipeline with a real +``run_shield_moderation_v2`` / ``build_shield`` → ``GraniteGuardian.run()`` +path for ``/responses`` and ``/rlsapi``, and the pydantic-ai capability +path via ``build_agent`` → ``GraniteGuardian.wrap_run()`` for ``/query`` +and ``/streaming_query``. + +The Guardian's LLM call (``model_request``) is mocked so no real +inference server is needed. +""" + +# pylint: disable=too-many-arguments +# pylint: disable=too-many-positional-arguments + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import pytest +from fastapi import Request +from fastapi.responses import StreamingResponse +from ogx_client.models.open_ai_response_object import OpenAIResponseObject +from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart +from pydantic_ai.models import ModelRequestParameters, StreamedResponse +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.settings import ModelSettings +from pydantic_ai.usage import RequestUsage +from pytest_mock import MockerFixture +from sqlalchemy.orm import Session + +from app.endpoints.query import query_endpoint_handler +from app.endpoints.responses import responses_endpoint_handler +from app.endpoints.rlsapi_v1 import infer_endpoint +from app.endpoints.streaming_query import streaming_query_endpoint_handler +from authentication.interface import AuthTuple +from configuration import AppConfig +from models.api.requests import QueryRequest, ResponsesRequest +from models.api.requests.rlsapi import RlsapiV1InferRequest +from models.api.responses.successful import ResponsesResponse +from models.api.responses.successful.rlsapi import RlsapiV1InferResponse +from models.common.responses.contexts import ResponsesContext +from models.config import ( + GraniteGuardianConfig, + GraniteGuardianShieldConfiguration, + RiskDefinition, +) +from models.database.conversations import UserConversation, UserTurn +from tests.integration.conftest import ( + make_openai_model, + make_openai_models_list_response, +) +from version import __version__ + +_GUARDIAN_MODULE = "pydantic_ai_lightspeed.capabilities.granite_guardian._capability" + +MOCK_CONV_ID = "conv_" + "a" * 48 +NORMALIZED_CONV_ID = "a" * 48 + +VIOLATION_MESSAGE = "Content blocked by Granite Guardian." + +_RESPONSE_DUMP: dict[str, Any] = { + "id": "resp_guardian_test", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "test-provider/test-model", + "store": False, + "output": [ + { + "type": "message", + "id": "msg-1", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Ansible is an automation tool.", + "annotations": [], + } + ], + } + ], + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} + + +def _guardian_shield_config() -> GraniteGuardianShieldConfiguration: + """Build a GraniteGuardianShieldConfiguration for testing.""" + return GraniteGuardianShieldConfiguration( + name="granite-guardian", + provider_id="granite_guardian", + config=GraniteGuardianConfig( + url="http://localhost:8080/v1", + risks=[ + RiskDefinition( + name="harmful_content", + description="Content that is harmful", + threshold=0.5, + points=["input"], + violation_message=VIOLATION_MESSAGE, + ) + ], + ), + ) + + +def _mock_guardian_model_request(mocker: MockerFixture) -> Any: + """Mock the Guardian's model_request to return valid logprobs. + + Returns the mock so callers can inspect call args. + """ + mock_response = mocker.Mock() + mock_response.usage = RequestUsage(input_tokens=5, output_tokens=1) + mock_response.provider_details = {"logprobs": [{"token": "no"}]} + return mocker.patch( + f"{_GUARDIAN_MODULE}.model_request", + new=mocker.AsyncMock(return_value=mock_response), + ) + + +def _mock_guardian_init(mocker: MockerFixture) -> None: + """Mock GraniteGuardian's external dependencies so __post_init__ succeeds.""" + mocker.patch(f"{_GUARDIAN_MODULE}.httpx.AsyncClient") + mocker.patch(f"{_GUARDIAN_MODULE}.AsyncOpenAI") + mocker.patch(f"{_GUARDIAN_MODULE}.OpenAIProvider") + mocker.patch(f"{_GUARDIAN_MODULE}.OpenAIChatModel") + + +# ============================================================================ +# /responses endpoint helpers +# ============================================================================ + + +def _build_responses_mock_client(mocker: MockerFixture) -> Any: + """Build a mock OGX client for responses integration tests.""" + mock_client = mocker.AsyncMock() + mock_client.responses.create = mocker.AsyncMock( + return_value=OpenAIResponseObject.from_dict(_RESPONSE_DUMP) + ) + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model() + ) + mock_client.shields.list.return_value = [] + mock_client.vector_stores.list.return_value = [] + mock_conv = mocker.MagicMock() + mock_conv.id = MOCK_CONV_ID + mock_client.conversations.create = mocker.AsyncMock(return_value=mock_conv) + return mock_client + + +def _patch_responses_client_holders(mocker: MockerFixture, mock_client: Any) -> None: + """Patch AsyncOgxClientHolder for the responses endpoint.""" + for module in ("app.endpoints.responses", "utils.endpoints"): + holder = mocker.patch(f"{module}.AsyncOgxClientHolder") + holder.return_value.get_client.return_value = mock_client + + original_cls = ResponsesContext + + def _skip_validation(**kwargs: Any) -> ResponsesContext: + return original_cls.model_construct(**kwargs) + + mocker.patch( + "app.endpoints.responses.ResponsesContext", side_effect=_skip_validation + ) + + +def _setup_responses_test(mocker: MockerFixture) -> Any: + """Set up mock client and patches for a responses integration test.""" + mock_client = _build_responses_mock_client(mocker) + _patch_responses_client_holders(mocker, mock_client) + mocker.patch( + "app.endpoints.responses.maybe_get_topic_summary", + new=mocker.AsyncMock(return_value=None), + ) + return mock_client + + +def _inject_guardian_shield(test_config: AppConfig) -> None: + """Add the Granite Guardian shield to the test configuration.""" + test_config.configuration.shields = [_guardian_shield_config()] + + +# ============================================================================ +# /responses endpoint tests +# ============================================================================ + + +class TestResponsesGraniteGuardian: + """Integration tests for Granite Guardian on the /responses endpoint.""" + + @pytest.mark.asyncio + async def test_blocks_unsafe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_db_session: Session, + test_auth: AuthTuple, + ) -> None: + """Test that Granite Guardian blocks unsafe input in non-streaming mode.""" + _, _ = test_config, test_db_session + _inject_guardian_shield(test_config) + mock_client = _setup_responses_test(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + request = ResponsesRequest( + input="Some harmful content", + model="test-provider/test-model", + stream=False, + store=True, + generate_topic_summary=False, + ) + + response = await responses_endpoint_handler( + request=test_request, + responses_request=request, + auth=test_auth, + mcp_headers={}, + ) + + assert isinstance(response, ResponsesResponse) + assert VIOLATION_MESSAGE in (response.output_text or "") + mock_client.responses.create.assert_not_called() + + # Check the block message shows up in the conversation turn + mock_client.items.create.assert_called_once() + items = mock_client.items.create.call_args[1]["add_items_request"].items + assert len(items) == 2 + + user_msg = items[0].actual_instance + assert user_msg.role == "user" + assert user_msg.content == "Some harmful content" + + assistant_msg = items[1].actual_instance + assert assistant_msg.role == "assistant" + assert assistant_msg.content == VIOLATION_MESSAGE + + @pytest.mark.asyncio + async def test_passes_safe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_db_session: Session, + test_auth: AuthTuple, + ) -> None: + """Test that Granite Guardian allows safe input through.""" + _, _ = test_config, test_db_session + _inject_guardian_shield(test_config) + mock_client = _setup_responses_test(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=True) + + request = ResponsesRequest( + input="What is Ansible?", + model="test-provider/test-model", + stream=False, + store=True, + generate_topic_summary=False, + ) + + response = await responses_endpoint_handler( + request=test_request, + responses_request=request, + auth=test_auth, + mcp_headers={}, + ) + + assert isinstance(response, ResponsesResponse) + assert response.id == "resp_guardian_test" + assert response.output_text == "Ansible is an automation tool." + mock_client.responses.create.assert_called_once() + + @pytest.mark.asyncio + async def test_streaming_blocks_unsafe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_db_session: Session, + test_auth: AuthTuple, + ) -> None: + """Test that Granite Guardian blocks unsafe input in streaming mode.""" + _, _ = test_config, test_db_session + _inject_guardian_shield(test_config) + mock_client = _setup_responses_test(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + request = ResponsesRequest( + input="Some harmful content", + model="test-provider/test-model", + stream=True, + store=True, + generate_topic_summary=False, + ) + + response = await responses_endpoint_handler( + request=test_request, + responses_request=request, + auth=test_auth, + mcp_headers={}, + ) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + + body = b"" + async for part in response.body_iterator: + if isinstance(part, str): + body += part.encode() + else: + body += bytes(part) + body_str = body.decode() + + assert "event: response.created" in body_str + assert "event: response.completed" in body_str + assert VIOLATION_MESSAGE in body_str + mock_client.responses.create.assert_not_called() + + @pytest.mark.asyncio + async def test_blocked_persists_turn( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_db_session: Session, + test_auth: AuthTuple, + ) -> None: + """Test that a blocked response persists the moderation turn to the DB.""" + _ = test_config + _inject_guardian_shield(test_config) + _setup_responses_test(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + request = ResponsesRequest( + input="Blocked content", + model="test-provider/test-model", + stream=False, + store=True, + generate_topic_summary=False, + ) + + await responses_endpoint_handler( + request=test_request, + responses_request=request, + auth=test_auth, + mcp_headers={}, + ) + + conversation = ( + test_db_session.query(UserConversation) + .filter_by(id=NORMALIZED_CONV_ID) + .first() + ) + assert conversation is not None + assert conversation.last_response_id is None + + turns = ( + test_db_session.query(UserTurn) + .filter_by(conversation_id=NORMALIZED_CONV_ID) + .all() + ) + assert len(turns) == 1 + assert turns[0].response_id.startswith("modr-") + + +# ============================================================================ +# /rlsapi endpoint helpers +# ============================================================================ + + +def _create_rlsapi_mock_request(mocker: MockerFixture) -> Any: + """Create a mock FastAPI Request for rlsapi tests.""" + mock_request = mocker.Mock() + mock_request.state = mocker.Mock(spec=[]) + mock_request.headers = {"User-Agent": f"CLA/{__version__}"} + return mock_request + + +def _setup_rlsapi_responses_mock(mocker: MockerFixture) -> Any: + """Set up responses.create mock for rlsapi tests.""" + mock_response = mocker.Mock() + mock_output = mocker.Mock() + mock_output.type = "message" + mock_output.role = "assistant" + mock_output.content = "Use the `ls` command to list files." + mock_response.output = [mock_output] + mock_usage = mocker.Mock() + mock_usage.input_tokens = 10 + mock_usage.output_tokens = 5 + mock_response.usage = mock_usage + + mock_responses = mocker.Mock() + mock_responses.create = mocker.AsyncMock(return_value=mock_response) + mock_client = mocker.Mock() + mock_client.responses = mock_responses + + mock_holder_class = mocker.patch("app.endpoints.rlsapi_v1.AsyncOgxClientHolder") + mock_holder_class.return_value.get_client.return_value = mock_client + return mock_client + + +# ============================================================================ +# /rlsapi endpoint tests +# ============================================================================ + + +class TestRlsapiGraniteGuardian: + """Integration tests for Granite Guardian on the /rlsapi v1 /infer endpoint.""" + + @pytest.fixture(name="rlsapi_config") + def rlsapi_config_fixture( + self, test_config: AppConfig, mocker: MockerFixture + ) -> AppConfig: + """Extend test_config with rlsapi defaults and Granite Guardian shield.""" + test_config.inference.default_model = "test-model" + test_config.inference.default_provider = "test-provider" + test_config.configuration.shields = [_guardian_shield_config()] + mocker.patch("app.endpoints.rlsapi_v1.configuration", test_config) + return test_config + + @pytest.fixture(name="mock_model_configured") + def mock_model_configured_fixture(self, mocker: MockerFixture) -> None: + """Mock model existence check to pass.""" + mocker.patch( + "app.endpoints.rlsapi_v1.check_model_configured", + new=mocker.AsyncMock(return_value=True), + ) + + @pytest.mark.asyncio + async def test_blocks_unsafe_input( + self, + rlsapi_config: AppConfig, + mock_model_configured: None, + mocker: MockerFixture, + test_auth: AuthTuple, + ) -> None: + """Test that Granite Guardian blocks unsafe input on /infer.""" + _, _ = rlsapi_config, mock_model_configured + mock_client = _setup_rlsapi_responses_mock(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + response = await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="Harmful question"), + request=_create_rlsapi_mock_request(mocker), + background_tasks=mocker.Mock(), + auth=test_auth, + ) + + assert isinstance(response, RlsapiV1InferResponse) + assert response.data.text == VIOLATION_MESSAGE + mock_client.responses.create.assert_not_called() + + @pytest.mark.asyncio + async def test_passes_safe_input( + self, + rlsapi_config: AppConfig, + mock_model_configured: None, + mocker: MockerFixture, + test_auth: AuthTuple, + ) -> None: + """Test that Granite Guardian allows safe input on /infer.""" + _, _ = rlsapi_config, mock_model_configured + mock_client = _setup_rlsapi_responses_mock(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=True) + + response = await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="How do I list files?"), + request=_create_rlsapi_mock_request(mocker), + background_tasks=mocker.Mock(), + auth=test_auth, + ) + + assert isinstance(response, RlsapiV1InferResponse) + assert response.data.text == "Use the `ls` command to list files." + mock_client.responses.create.assert_called_once() + + @pytest.mark.asyncio + async def test_blocked_response_has_no_token_counts( + self, + rlsapi_config: AppConfig, + mock_model_configured: None, + mocker: MockerFixture, + test_auth: AuthTuple, + ) -> None: + """Test that a blocked rlsapi response has no token usage counts.""" + _, _ = rlsapi_config, mock_model_configured + _setup_rlsapi_responses_mock(mocker) + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + response = await infer_endpoint( + infer_request=RlsapiV1InferRequest(question="Bad input"), + request=_create_rlsapi_mock_request(mocker), + background_tasks=mocker.Mock(), + auth=test_auth, + ) + + assert response.data.input_tokens is None + assert response.data.output_tokens is None + + +# ============================================================================ +# /query and /streaming_query shared helpers +# ============================================================================ + + +# Signature required by pydantic-ai FunctionModel handler interface. +async def _mock_llm(messages: list[Any], info: Any) -> ModelResponse: + """FunctionModel handler returning a simple assistant response.""" + _, _ = messages, info + return ModelResponse( + parts=[TextPart("This is a test response about Ansible.")], + finish_reason="stop", + provider_response_id="response-123", + ) + + +# Signature required by pydantic-ai FunctionModel stream handler interface. +async def _mock_llm_stream(messages: list[Any], info: Any) -> AsyncIterator[str]: + """FunctionModel stream handler yielding a simple assistant response.""" + _, _ = messages, info + yield "This is a test response about Ansible." + + +class _TestLLMModel(FunctionModel): + """FunctionModel subclass that sets finish_reason on streamed responses.""" + + @asynccontextmanager + async def request_stream( + self, + messages: list[ModelMessage], + model_settings: ModelSettings | None, + model_request_parameters: ModelRequestParameters, + run_context: Any = None, + ) -> AsyncIterator[StreamedResponse]: + """Delegate to parent and patch finish_reason on the streamed response.""" + async with super().request_stream( + messages, model_settings, model_request_parameters, run_context + ) as response: + response.finish_reason = "stop" + response.provider_response_id = "response-123" + yield response + + +def _mock_ogx_for_query(mocker: MockerFixture, module: str) -> Any: + """Patch AsyncOgxClientHolder in the given endpoint module and return the mock client.""" + mock_holder_class = mocker.patch(f"{module}.AsyncOgxClientHolder") + mock_client = mocker.AsyncMock() + + mock_client.openai.list.return_value = make_openai_models_list_response( + make_openai_model() + ) + mock_client.vector_stores.list.return_value = [] + mock_client.shields.list.return_value = [] + mock_client.items.create = mocker.AsyncMock() + + mock_conv = mocker.MagicMock() + mock_conv.id = MOCK_CONV_ID + mock_client.conversations.create = mocker.AsyncMock(return_value=mock_conv) + + mock_holder_class.return_value.get_client.return_value = mock_client + return mock_client + + +def _mock_build_agent_model(mocker: MockerFixture) -> None: + """Replace OgxResponsesModel.from_ogx_client with a test model.""" + mocker.patch( + "utils.pydantic_ai_helpers.OgxResponsesModel.from_ogx_client", + return_value=_TestLLMModel( + _mock_llm, + stream_function=_mock_llm_stream, + settings={"extra_body": {"conversation": MOCK_CONV_ID}}, + ), + ) + + +def _mock_guardian_for_capability(mocker: MockerFixture) -> Any: + """Mock all Guardian external dependencies for the capability path. + + Includes AsyncOgxClientHolder in the Guardian module so + ``wrap_run`` can call ``append_turn_to_conversation``. + + Returns: + The mock OGX client used by the Guardian capability. + """ + _mock_guardian_init(mocker) + _mock_guardian_model_request(mocker) + mock_holder = mocker.patch(f"{_GUARDIAN_MODULE}.AsyncOgxClientHolder") + mock_client = mocker.AsyncMock() + mock_holder.return_value.get_client.return_value = mock_client + return mock_client + + +# ============================================================================ +# /query endpoint tests +# ============================================================================ + + +class TestQueryGraniteGuardian: + """Integration tests for Granite Guardian on the /query endpoint. + + These tests let ``build_agent`` run for real so ``GraniteGuardian`` + is wired as a pydantic-ai capability and ``wrap_run`` fires during + ``agent.run()``. The main LLM model is replaced by a + ``FunctionModel`` and the Guardian's ``model_request`` is mocked. + """ + + @pytest.mark.asyncio + async def test_blocks_unsafe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_auth: AuthTuple, + ) -> None: + """Test that Guardian capability blocks unsafe input on /query.""" + _inject_guardian_shield(test_config) + _mock_ogx_for_query(mocker, "app.endpoints.query") + _mock_build_agent_model(mocker) + guardian_client = _mock_guardian_for_capability(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + response = await query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="Some harmful content"), + auth=test_auth, + mcp_headers={}, + ) + + assert VIOLATION_MESSAGE in response.response + + guardian_client.items.create.assert_called_once() + items = guardian_client.items.create.call_args[1]["add_items_request"].items + assert len(items) == 2 + + user_msg = items[0].actual_instance + assert user_msg.role == "user" + assert user_msg.content == "Some harmful content" + + assistant_msg = items[1].actual_instance + assert assistant_msg.role == "assistant" + assert assistant_msg.content == VIOLATION_MESSAGE + + @pytest.mark.asyncio + async def test_passes_safe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_auth: AuthTuple, + ) -> None: + """Test that Guardian capability allows safe input on /query.""" + _inject_guardian_shield(test_config) + _mock_ogx_for_query(mocker, "app.endpoints.query") + _mock_build_agent_model(mocker) + _mock_guardian_for_capability(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=True) + + response = await query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="What is Ansible?"), + auth=test_auth, + mcp_headers={}, + ) + + assert "Ansible" in response.response + assert VIOLATION_MESSAGE not in response.response + + @pytest.mark.asyncio + async def test_guardian_token_usage_is_not_included_in_user_token_usage( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_auth: AuthTuple, + ) -> None: + """Test that a blocked /query response reports only Guardian token usage.""" + _inject_guardian_shield(test_config) + _mock_ogx_for_query(mocker, "app.endpoints.query") + _mock_build_agent_model(mocker) + _mock_guardian_for_capability(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", side_effect=[False, True]) + + response = await query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="Bad content"), + auth=test_auth, + mcp_headers={}, + ) + + assert VIOLATION_MESSAGE in response.response + assert response.input_tokens == 0 + assert response.output_tokens == 0 + + response = await query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="Bad content"), + auth=test_auth, + mcp_headers={}, + ) + + assert VIOLATION_MESSAGE not in response.response + # Both token counts are rough estimates produced by FunctionModel's + # _estimate_usage: input = 50-token base + prompt tokens, output = response text tokens. + assert response.input_tokens == 52 + assert response.output_tokens == 8 + + +# ============================================================================ +# /streaming_query endpoint tests +# ============================================================================ + + +class TestStreamingQueryGraniteGuardian: + """Integration tests for Granite Guardian on the /streaming_query endpoint. + + Uses the same capability path as ``/query`` but exercises + ``agent.run_stream_events()`` instead of ``agent.run()``. + """ + + @pytest.mark.asyncio + async def test_blocks_unsafe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_auth: AuthTuple, + ) -> None: + """Test that Guardian capability blocks unsafe input on /streaming_query.""" + _inject_guardian_shield(test_config) + _mock_ogx_for_query(mocker, "app.endpoints.streaming_query") + _mock_build_agent_model(mocker) + guardian_client = _mock_guardian_for_capability(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=False) + + response = await streaming_query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="Some harmful content"), + auth=test_auth, + mcp_headers={}, + ) + + assert isinstance(response, StreamingResponse) + + body = b"" + async for part in response.body_iterator: + if isinstance(part, str): + body += part.encode() + else: + body += bytes(part) + body_str = body.decode() + + assert VIOLATION_MESSAGE in body_str + + guardian_client.items.create.assert_called_once() + items = guardian_client.items.create.call_args[1]["add_items_request"].items + assert len(items) == 2 + + user_msg = items[0].actual_instance + assert user_msg.role == "user" + assert user_msg.content == "Some harmful content" + + assistant_msg = items[1].actual_instance + assert assistant_msg.role == "assistant" + assert assistant_msg.content == VIOLATION_MESSAGE + + @pytest.mark.asyncio + async def test_passes_safe_input( + self, + test_config: AppConfig, + mocker: MockerFixture, + test_request: Request, + test_auth: AuthTuple, + ) -> None: + """Test that Guardian capability allows safe input on /streaming_query.""" + _inject_guardian_shield(test_config) + _mock_ogx_for_query(mocker, "app.endpoints.streaming_query") + _mock_build_agent_model(mocker) + _mock_guardian_for_capability(mocker) + mocker.patch(f"{_GUARDIAN_MODULE}.is_safe", return_value=True) + + response = await streaming_query_endpoint_handler( + request=test_request, + query_request=QueryRequest(query="What is Ansible?"), + auth=test_auth, + mcp_headers={}, + ) + + assert isinstance(response, StreamingResponse) + + body = b"" + async for part in response.body_iterator: + if isinstance(part, str): + body += part.encode() + else: + body += bytes(part) + body_str = body.decode() + + assert "Ansible" in body_str + assert VIOLATION_MESSAGE not in body_str diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py new file mode 100644 index 000000000..799a2b050 --- /dev/null +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/__init__.py @@ -0,0 +1 @@ +"""Unit tests for Granite Guardian capability.""" diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_capability.py b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_capability.py new file mode 100644 index 000000000..16bb07d22 --- /dev/null +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_capability.py @@ -0,0 +1,669 @@ +"""Unit tests for pydantic_ai_lightspeed.capabilities.granite_guardian._capability module.""" + +# pylint: disable=protected-access +# pylint: disable=too-many-arguments +# pylint: disable=too-many-positional-arguments + +import pytest +from pydantic_ai import AgentRunResult, RunContext +from pydantic_ai.exceptions import ModelAPIError, UnexpectedModelBehavior +from pydantic_ai.usage import RequestUsage, RunUsage +from pytest_mock import MockerFixture, MockType + +from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.config import GraniteGuardianConfig, GuardrailPoint, RiskDefinition +from pydantic_ai_lightspeed.capabilities.granite_guardian._capability import ( + GraniteGuardian, + _filter_guardrails, + _get_batch_size, + _package_risk_check_task, + _run_risk_check, +) + +_MODULE = "pydantic_ai_lightspeed.capabilities.granite_guardian._capability" + + +def _make_risk( + name: str = "test_risk", + description: str = "test description", + threshold: float = 0.5, + points: list[GuardrailPoint] | None = None, + enabled: bool = True, + violation_message: str = "Blocked.", +) -> RiskDefinition: + """Build a RiskDefinition for testing.""" + return RiskDefinition( + name=name, + description=description, + threshold=threshold, + points=points or ["input"], + enabled=enabled, + violation_message=violation_message, + ) + + +def _make_config( + url: str = "https://example.com/v1", + risks: list[RiskDefinition] | None = None, + api_key: str | None = None, + model_id: str | None = None, +) -> GraniteGuardianConfig: + """Build a GraniteGuardianConfig for testing.""" + kwargs: dict = { + "url": url, + "api_key": api_key, + "risks": risks or [_make_risk()], + } + if model_id is not None: + kwargs["model_id"] = model_id + return GraniteGuardianConfig(**kwargs) + + +class TestGetBatchSize: + """Tests for _get_batch_size parallel-to-batch-size resolution.""" + + def test_true_returns_num_guardrails(self) -> None: + """Test that True resolves to the total number of guardrails.""" + assert _get_batch_size(True, 5) == 5 + + def test_true_with_zero_guardrails_returns_one(self) -> None: + """Test that True with zero guardrails returns 1 to avoid range step=0.""" + assert _get_batch_size(True, 0) == 1 + + def test_false_returns_one(self) -> None: + """Test that False resolves to sequential execution (batch size 1).""" + assert _get_batch_size(False, 5) == 1 + + def test_int_returns_as_is(self) -> None: + """Test that an explicit integer is returned unchanged.""" + assert _get_batch_size(3, 10) == 3 + + +class TestFilterGuardrails: + """Tests for _filter_guardrails.""" + + def test_filters_by_point(self) -> None: + """Test that only risks matching the point are returned.""" + risks = [ + _make_risk(name="input_only", points=["input"]), + _make_risk(name="output_only", points=["output"]), + _make_risk(name="both", points=["input", "output"]), + ] + result = _filter_guardrails(risks, "input") + names = [g[0] for g in result] + assert names == ["input_only", "both"] + + def test_filters_disabled_risks(self) -> None: + """Test that disabled risks are excluded.""" + risks = [ + _make_risk(name="enabled", enabled=True), + _make_risk(name="disabled", enabled=False), + ] + result = _filter_guardrails(risks, "input") + assert len(result) == 1 + assert result[0][0] == "enabled" + + def test_empty_when_no_match(self) -> None: + """Test that an empty list is returned when no risks match.""" + risks = [_make_risk(points=["output"])] + result = _filter_guardrails(risks, "input") + assert result == [] + + def test_returns_correct_tuple_structure(self) -> None: + """Test that each guardrail tuple has the expected fields.""" + risks = [ + _make_risk( + name="harm", + description="harmful content", + threshold=0.7, + violation_message="Content blocked.", + ) + ] + result = _filter_guardrails(risks, "input") + assert len(result) == 1 + name, block, threshold, message = result[0] + assert name == "harm" + assert "harmful content" in block + assert threshold == 0.7 + assert message == "Content blocked." + + def test_thinking_mode_in_block(self) -> None: + """Test that enable_thinking produces a think-mode block.""" + risk = _make_risk() + risk.enable_thinking = True + result = _filter_guardrails([risk], "input") + assert "" in result[0][1] + + +class TestRunRiskCheck: + """Tests for _run_risk_check.""" + + @pytest.mark.asyncio + async def test_returns_none_when_all_safe(self, mocker: MockerFixture) -> None: + """Test that None is returned when all guardrails pass.""" + mocker.patch(f"{_MODULE}.is_safe", return_value=True) + mock_response = mocker.Mock() + mock_response.usage = RequestUsage(input_tokens=10, output_tokens=1) + mock_response.provider_details = {"logprobs": [{"token": "no"}]} + mocker.patch(f"{_MODULE}.model_request", return_value=mock_response) + + guardrails = [("risk1", "block1", 0.5, "Blocked 1")] + violation, usage = await _run_risk_check("hello", mocker.Mock(), guardrails) + + assert violation is None + assert usage.input_tokens == 10 + + @pytest.mark.asyncio + async def test_returns_violation_on_first_failure( + self, mocker: MockerFixture + ) -> None: + """Test that the first violated guardrail's message is returned.""" + mocker.patch(f"{_MODULE}.is_safe", side_effect=[True, False]) + mock_response = mocker.Mock() + mock_response.usage = RequestUsage(input_tokens=5, output_tokens=1) + mock_response.provider_details = {"logprobs": [{"token": "yes"}]} + mocker.patch(f"{_MODULE}.model_request", return_value=mock_response) + + guardrails = [ + ("risk1", "block1", 0.5, "Blocked 1"), + ("risk2", "block2", 0.5, "Blocked 2"), + ] + violation, usage = await _run_risk_check("hello", mocker.Mock(), guardrails) + + assert violation == "Blocked 2" + assert usage.input_tokens == 10 + + @pytest.mark.asyncio + async def test_returns_empty_usage_for_no_guardrails( + self, mocker: MockerFixture + ) -> None: + """Test that empty guardrails return None with zero usage.""" + violation, usage = await _run_risk_check("hello", mocker.Mock(), []) + assert violation is None + assert usage.input_tokens == 0 + + @pytest.mark.asyncio + async def test_raises_when_no_provider_details(self, mocker: MockerFixture) -> None: + """Test that missing provider_details raises UnexpectedModelBehavior.""" + mock_response = mocker.Mock() + mock_response.usage = RequestUsage() + mock_response.provider_details = None + mocker.patch(f"{_MODULE}.model_request", return_value=mock_response) + + guardrails = [("risk1", "block1", 0.5, "Blocked")] + with pytest.raises(UnexpectedModelBehavior, match="No provider_details"): + await _run_risk_check("hello", mocker.Mock(), guardrails) + + @pytest.mark.asyncio + async def test_raises_when_no_logprobs(self, mocker: MockerFixture) -> None: + """Test that missing logprobs raises UnexpectedModelBehavior.""" + mock_response = mocker.Mock() + mock_response.usage = RequestUsage() + mock_response.provider_details = {"logprobs": None} + mocker.patch(f"{_MODULE}.model_request", return_value=mock_response) + + guardrails = [("risk1", "block1", 0.5, "Blocked")] + with pytest.raises(UnexpectedModelBehavior, match="No logprobs"): + await _run_risk_check("hello", mocker.Mock(), guardrails) + + +class TestRunRiskCheckBatching: + """Tests for _run_risk_check batch parallelism and cross-batch short-circuiting.""" + + def _mock_response(self, mocker: MockerFixture) -> MockType: + """Build a mock model response with valid provider_details.""" + resp = mocker.Mock() + resp.usage = RequestUsage(input_tokens=5, output_tokens=1) + resp.provider_details = {"logprobs": [{"token": "no"}]} + return resp + + @pytest.mark.asyncio + async def test_violation_in_first_batch_skips_second_batch( + self, mocker: MockerFixture + ) -> None: + """Test that a violation in batch 1 prevents batch 2 from executing.""" + mock_model_request = mocker.patch( + f"{_MODULE}.model_request", return_value=self._mock_response(mocker) + ) + mocker.patch(f"{_MODULE}.is_safe", return_value=False) + + guardrails = [(f"risk{i}", f"block{i}", 0.5, f"Blocked {i}") for i in range(5)] + violation, _ = await _run_risk_check("hello", mocker.Mock(), guardrails) + + assert violation == "Blocked 0" + assert mock_model_request.call_count == 3 + + @pytest.mark.asyncio + async def test_all_batches_run_when_no_violation( + self, mocker: MockerFixture + ) -> None: + """Test that all guardrails across batches are checked when safe.""" + mock_model_request = mocker.patch( + f"{_MODULE}.model_request", return_value=self._mock_response(mocker) + ) + mocker.patch(f"{_MODULE}.is_safe", return_value=True) + + guardrails = [(f"risk{i}", f"block{i}", 0.5, f"Blocked {i}") for i in range(5)] + violation, usage = await _run_risk_check("hello", mocker.Mock(), guardrails) + + assert violation is None + assert mock_model_request.call_count == 5 + assert usage.input_tokens == 25 + + @pytest.mark.asyncio + async def test_token_usage_includes_violating_batch( + self, mocker: MockerFixture + ) -> None: + """Test that token usage from the batch containing the violation is accumulated.""" + mocker.patch( + f"{_MODULE}.model_request", return_value=self._mock_response(mocker) + ) + mocker.patch(f"{_MODULE}.is_safe", side_effect=[True, True, False]) + + guardrails = [(f"risk{i}", f"block{i}", 0.5, f"Blocked {i}") for i in range(3)] + _, usage = await _run_risk_check("hello", mocker.Mock(), guardrails) + + assert usage.input_tokens == 15 + + @pytest.mark.asyncio + async def test_logs_risk_name_and_latency(self, mocker: MockerFixture) -> None: + """Test that each risk check logs its name and elapsed time at INFO level.""" + mock_response = mocker.Mock() + mock_response.usage = RequestUsage() + mocker.patch(f"{_MODULE}.model_request", return_value=mock_response) + mock_logger = mocker.patch(f"{_MODULE}.logger") + + guardrail = ("harm_detection", "block", 0.5, "Blocked") + await _package_risk_check_task("hello", guardrail, mocker.Mock()) + + mock_logger.info.assert_called_once() + log_args = mock_logger.info.call_args + assert "harm_detection" in log_args[0][1] + + +class TestGraniteGuardianInit: + """Tests for GraniteGuardian initialization.""" + + @pytest.fixture(autouse=True) + def _mock_init(self, mocker: MockerFixture) -> None: + """Mock model creation and clear cache for all tests.""" + GraniteGuardian._model_cache.clear() + mocker.patch(f"{_MODULE}.httpx.AsyncClient") + mocker.patch(f"{_MODULE}.AsyncOpenAI") + mocker.patch(f"{_MODULE}.OpenAIProvider") + mocker.patch(f"{_MODULE}.OpenAIChatModel") + + def test_creates_model_on_init(self, mocker: MockerFixture) -> None: + """Test that __post_init__ creates the OpenAI model.""" + mock_provider = mocker.patch(f"{_MODULE}.OpenAIProvider") + mock_model = mocker.patch(f"{_MODULE}.OpenAIChatModel") + + config = _make_config() + guardian = GraniteGuardian(config=config) + + mock_provider.assert_called_once() + mock_model.assert_called_once() + assert guardian._model is not None + + def test_custom_model_name_passed_to_chat_model( + self, mocker: MockerFixture + ) -> None: + """Test that a custom model name from config is forwarded to OpenAIChatModel.""" + mock_model = mocker.patch(f"{_MODULE}.OpenAIChatModel") + + config = _make_config(model_id="custom/guardian-local") + GraniteGuardian(config=config) + + mock_model.assert_called_once() + assert mock_model.call_args[0][0] == "custom/guardian-local" + + def test_api_key_passed_to_client(self, mocker: MockerFixture) -> None: + """Test that the API key is extracted and passed to the OpenAI client.""" + mock_openai = mocker.patch(f"{_MODULE}.AsyncOpenAI") + + config = _make_config(risks=[_make_risk()], api_key="test-key") + GraniteGuardian(config=config) + + _, kwargs = mock_openai.call_args + assert kwargs["api_key"] == "test-key" + + def test_no_api_key_passes_fallback(self, mocker: MockerFixture) -> None: + """Test that None api_key passes fallback value to the OpenAI client.""" + mock_openai = mocker.patch(f"{_MODULE}.AsyncOpenAI") + + config = _make_config() + GraniteGuardian(config=config) + + _, kwargs = mock_openai.call_args + assert kwargs["api_key"] == "api-key-not-set" + + def test_raises_when_api_key_with_non_https_url(self) -> None: + """Test that a ValueError is raised when api_key is set but URL is not HTTPS.""" + config = _make_config(url="http://example.com/v1", api_key="test-key") + with pytest.raises( + ValueError, + match="Granite Guardian endpoints with an API key must use HTTPS", + ): + GraniteGuardian(config=config) + + def test_reuses_cached_model_for_same_config(self) -> None: + """Test that instances sharing a config object reuse the cached model.""" + config = _make_config() + guardian_a = GraniteGuardian(config=config) + guardian_b = GraniteGuardian(config=config) + + assert guardian_a._model is guardian_b._model + + def test_different_configs_get_different_models( + self, mocker: MockerFixture + ) -> None: + """Test that different config objects produce separate cached models.""" + mock_model = mocker.patch(f"{_MODULE}.OpenAIChatModel") + + config_a = _make_config() + config_b = _make_config() + GraniteGuardian(config=config_a) + GraniteGuardian(config=config_b) + + assert mock_model.call_count == 2 + + +class TestGraniteGuardianWrapRun: + """Tests for GraniteGuardian.wrap_run method.""" + + @pytest.fixture(autouse=True) + def _mock_init(self, mocker: MockerFixture) -> None: + """Mock model creation and clear cache for all tests.""" + GraniteGuardian._model_cache.clear() + mocker.patch(f"{_MODULE}.httpx.AsyncClient") + mocker.patch(f"{_MODULE}.AsyncOpenAI") + mocker.patch(f"{_MODULE}.OpenAIProvider") + mocker.patch(f"{_MODULE}.OpenAIChatModel") + mocker.patch(f"{_MODULE}.AsyncOgxClientHolder") + + @pytest.fixture(name="mock_append_turn", autouse=True) + def mock_append_turn_fixture(self, mocker: MockerFixture) -> MockType: + """Mock the conversation-persistence call used on rejection.""" + return mocker.patch( + f"{_MODULE}.append_turn_to_conversation", new_callable=mocker.AsyncMock + ) + + @pytest.fixture(name="mock_ctx") + def mock_ctx_fixture(self, mocker: MockerFixture) -> RunContext: + """Create a mock RunContext with a conversation ID.""" + ctx = mocker.Mock(spec=RunContext) + ctx.prompt = "How do I create a pod?" + ctx.usage = RunUsage() + ctx.model = mocker.Mock() + ctx.model.settings = {"extra_body": {"conversation": "conv_test"}} + return ctx + + @pytest.fixture(name="mock_handler") + def mock_handler_fixture(self, mocker: MockerFixture) -> MockType: + """Create a mock WrapRunHandler.""" + handler = mocker.AsyncMock() + handler.return_value = mocker.Mock(spec=AgentRunResult) + return handler + + @pytest.mark.asyncio + async def test_safe_input_calls_handler( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + ) -> None: + """Test that a safe input proceeds to the handler.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=(None, RequestUsage(input_tokens=5, output_tokens=1)), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_handler.assert_awaited_once() + assert result == mock_handler.return_value + + @pytest.mark.asyncio + async def test_violation_short_circuits( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + ) -> None: + """Test that a violation short-circuits without calling the handler.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=( + "Content blocked.", + RequestUsage(input_tokens=5, output_tokens=1), + ), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_handler.assert_not_awaited() + assert isinstance(result, AgentRunResult) + assert result.output == "Content blocked." + + @pytest.mark.asyncio + async def test_violation_persists_turn_to_conversation( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that a violation appends the turn to the conversation.""" + mock_client = mocker.Mock() + mocker.patch( + f"{_MODULE}.AsyncOgxClientHolder" + ).return_value.get_client.return_value = mock_client + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=( + "Content blocked.", + RequestUsage(input_tokens=5, output_tokens=1), + ), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_awaited_once_with( + mock_client, + "conv_test", + "How do I create a pod?", + "Content blocked.", + ) + + @pytest.mark.asyncio + async def test_violation_skips_persistence_when_no_conversation_id( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that persistence is skipped without a conversation ID.""" + mock_ctx.model = mocker.Mock(settings={}) + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=( + "Content blocked.", + RequestUsage(input_tokens=5, output_tokens=1), + ), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_not_awaited() + assert result.output == "Content blocked." + + @pytest.mark.asyncio + async def test_safe_input_does_not_persist_turn( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + mock_append_turn: MockType, + ) -> None: + """Test that a safe input does not touch the conversation.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=(None, RequestUsage()), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_append_turn.assert_not_awaited() + + @pytest.mark.asyncio + async def test_usage_is_not_accumulated( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + ) -> None: + """Test that guardian token usage is not added to the run context.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=(None, RequestUsage(input_tokens=20, output_tokens=5)), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + await guardian.wrap_run(mock_ctx, handler=mock_handler) + + assert mock_ctx.usage.input_tokens == 0 + assert mock_ctx.usage.output_tokens == 0 + + @pytest.mark.asyncio + async def test_risk_check_failure_propagates( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + ) -> None: + """Test that wrap_run does not swallow errors from the Guardian endpoint.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + side_effect=UnexpectedModelBehavior( + "No provider_details provided from granite guardian's response" + ), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + + with pytest.raises(UnexpectedModelBehavior, match="No provider_details"): + await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_handler.assert_not_awaited() + + @pytest.mark.asyncio + async def test_guardian_unreachable_propagates( + self, + mocker: MockerFixture, + mock_ctx: RunContext, + mock_handler: MockType, + ) -> None: + """Test that wrap_run does not swallow connection errors from the Guardian endpoint.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + side_effect=ModelAPIError("test", "Connection refused"), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + + with pytest.raises(ModelAPIError, match="Connection refused"): + await guardian.wrap_run(mock_ctx, handler=mock_handler) + + mock_handler.assert_not_awaited() + + +class TestGraniteGuardianRun: + """Tests for GraniteGuardian.run (standalone shield interface).""" + + @pytest.fixture(autouse=True) + def _mock_init(self, mocker: MockerFixture) -> None: + """Mock model creation and clear cache for all tests.""" + GraniteGuardian._model_cache.clear() + mocker.patch(f"{_MODULE}.httpx.AsyncClient") + mocker.patch(f"{_MODULE}.AsyncOpenAI") + mocker.patch(f"{_MODULE}.OpenAIProvider") + mocker.patch(f"{_MODULE}.OpenAIChatModel") + + @pytest.mark.asyncio + async def test_returns_passed_when_safe(self, mocker: MockerFixture) -> None: + """Test that a safe input returns ShieldModerationPassed.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=(None, RequestUsage()), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.run("safe text") + + assert isinstance(result, ShieldModerationPassed) + + @pytest.mark.asyncio + async def test_returns_blocked_on_violation(self, mocker: MockerFixture) -> None: + """Test that a violation returns ShieldModerationBlocked.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=("Content blocked.", RequestUsage()), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.run("harmful text") + + assert isinstance(result, ShieldModerationBlocked) + assert result.message == "Content blocked." + + @pytest.mark.asyncio + async def test_uses_run_moderation_guardrail_point( + self, mocker: MockerFixture + ) -> None: + """Test that run() filters by run_moderation_guardrail_point.""" + mock_filter = mocker.patch(f"{_MODULE}._filter_guardrails", return_value=[]) + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=(None, RequestUsage()), + ) + + config = _make_config() + guardian = GraniteGuardian( + config=config, run_moderation_guardrail_point="output" + ) + await guardian.run("some text") + + mock_filter.assert_called_once_with(config.risks, "output") + + @pytest.mark.asyncio + async def test_blocked_result_has_moderation_id( + self, mocker: MockerFixture + ) -> None: + """Test that blocked results include a moderation ID.""" + mocker.patch( + f"{_MODULE}._run_risk_check", + return_value=("Blocked.", RequestUsage()), + ) + + config = _make_config() + guardian = GraniteGuardian(config=config) + result = await guardian.run("bad text") + + assert isinstance(result, ShieldModerationBlocked) + assert result.moderation_id.startswith("modr-") diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_utils.py b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_utils.py new file mode 100644 index 000000000..72365835b --- /dev/null +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/granite_guardian/test_utils.py @@ -0,0 +1,344 @@ +"""Unit tests for pydantic_ai_lightspeed.capabilities.granite_guardian.utils module.""" + +import math + +import pytest +from openai.types.chat.chat_completion_token_logprob import ( + ChatCompletionTokenLogprob, + TopLogprob, +) +from pydantic_ai.exceptions import UnexpectedModelBehavior + +from pydantic_ai_lightspeed.capabilities.granite_guardian.utils import ( + _clean_up_candidates, + _extract_tokens_inside_score_tag, + _get_risky_probabilities, + _search_tag, + build_guardian_block, + is_safe, +) + + +def _make_logprob(token: str, logprob: float = 0.0, top_logprobs=None): + """Build a raw logprob dict matching the ChatCompletionTokenLogprob schema.""" + return { + "token": token, + "logprob": logprob, + "bytes": None, + "top_logprobs": top_logprobs or [], + } + + +def _make_token_logprob(token: str, logprob: float = 0.0, top_logprobs=None): + """Build a ChatCompletionTokenLogprob instance.""" + tops = [ + TopLogprob(token=t, logprob=lp, bytes=None) for t, lp in (top_logprobs or []) + ] + return ChatCompletionTokenLogprob( + token=token, logprob=logprob, bytes=None, top_logprobs=tops + ) + + +class TestBuildGuardianBlock: + """Tests for build_guardian_block.""" + + def test_nothink_mode_contains_no_think_tag(self) -> None: + """Test that no-think mode uses the no-think preamble.""" + result = build_guardian_block("test criteria", think=False) + assert "" in result + assert "### Criteria: test criteria" in result + assert "### Scoring Schema:" in result + + def test_think_mode_contains_think_tag(self) -> None: + """Test that think mode uses the think preamble.""" + result = build_guardian_block("test criteria", think=True) + assert "" in result + assert "" not in result + + def test_criteria_embedded_in_output(self) -> None: + """Test that the criteria text appears in the output.""" + result = build_guardian_block("harmful content detection") + assert "### Criteria: harmful content detection" in result + + def test_default_is_nothink(self) -> None: + """Test that think defaults to False.""" + result = build_guardian_block("criteria") + assert "" in result + + +class TestSearchTag: + """Tests for _search_tag.""" + + def test_finds_complete_tag(self) -> None: + """Test that a complete tag is detected.""" + remaining, found = _search_tag("", "") + assert found is True + assert remaining == "" + + def test_partial_tag_keeps_buffer(self) -> None: + """Test that a partial tag keeps the buffer intact.""" + remaining, found = _search_tag("", " None: + """Test that non-tag content clears the buffer.""" + remaining, found = _search_tag("", "hello") + assert found is False + assert remaining == "" + + def test_tag_embedded_in_longer_string(self) -> None: + """Test detection when tag is part of a longer string.""" + remaining, found = _search_tag("", "some content") + assert found is True + assert remaining == "" + + def test_empty_buffer(self) -> None: + """Test empty buffer is treated as non-tag content.""" + remaining, found = _search_tag("", "") + assert found is False + assert remaining == "" + + def test_whitespace_before_tag(self) -> None: + """Test that leading whitespace before a tag is accepted.""" + remaining, found = _search_tag("", " ") + assert found is True + assert remaining == "" + + +class TestCleanUpCandidates: + """Tests for _clean_up_candidates.""" + + def test_removes_end_tag_tokens(self) -> None: + """Test that tokens forming the end tag are removed.""" + candidates = [ + _make_token_logprob("No"), + _make_token_logprob(""), + ] + _clean_up_candidates(candidates) + assert len(candidates) == 1 + assert candidates[0].token == "No" + + def test_single_end_tag_token(self) -> None: + """Test removal when end tag is a single token.""" + candidates = [ + _make_token_logprob("yes"), + _make_token_logprob(""), + ] + _clean_up_candidates(candidates) + assert len(candidates) == 1 + assert candidates[0].token == "yes" + + def test_empty_candidates(self) -> None: + """Test that empty list does not raise.""" + candidates = [] + _clean_up_candidates(candidates) + assert not candidates + + +class TestExtractTokensInsideScoreTag: + """Tests for _extract_tokens_inside_score_tag.""" + + def test_extracts_score_token_nothink(self) -> None: + """Test extraction from a no-think response format.""" + logprobs = [ + _make_logprob(""), + _make_logprob("\n"), + _make_logprob(""), + _make_logprob(""), + _make_logprob("No", top_logprobs=[]), + _make_logprob(""), + ] + result = _extract_tokens_inside_score_tag(logprobs) + assert result.token == "No" + + def test_extracts_score_token_with_think(self) -> None: + """Test extraction from a think-mode response with reasoning content.""" + logprobs = [ + _make_logprob(""), + _make_logprob("The user is asking about pods."), + _make_logprob(""), + _make_logprob(""), + _make_logprob("yes", top_logprobs=[]), + _make_logprob(""), + ] + result = _extract_tokens_inside_score_tag(logprobs) + assert result.token == "yes" + + def test_raises_when_no_tags_present(self) -> None: + """Test that missing tag structure raises UnexpectedModelBehavior.""" + logprobs = [ + _make_logprob("This is just plain text."), + ] + with pytest.raises(UnexpectedModelBehavior, match="did not generate"): + _extract_tokens_inside_score_tag(logprobs) + + def test_raises_when_score_tag_empty(self) -> None: + """Test that an empty score tag raises UnexpectedModelBehavior.""" + logprobs = [ + _make_logprob(""), + _make_logprob(""), + _make_logprob(""), + _make_logprob(""), + ] + with pytest.raises(UnexpectedModelBehavior, match="No token found"): + _extract_tokens_inside_score_tag(logprobs) + + def test_raises_when_multiple_tokens_in_score(self) -> None: + """Test that multiple content tokens in score raises UnexpectedModelBehavior.""" + logprobs = [ + _make_logprob(""), + _make_logprob(""), + _make_logprob(""), + _make_logprob("y"), + _make_logprob("es"), + _make_logprob(""), + ] + with pytest.raises(UnexpectedModelBehavior, match="More than one token"): + _extract_tokens_inside_score_tag(logprobs) + + def test_split_tags_across_tokens(self) -> None: + """Test extraction when tags are split across multiple tokens.""" + logprobs = [ + _make_logprob("<"), + _make_logprob("think>"), + _make_logprob(""), + _make_logprob("<"), + _make_logprob("score>"), + _make_logprob("no", top_logprobs=[]), + _make_logprob(""), + ] + result = _extract_tokens_inside_score_tag(logprobs) + assert result.token == "no" + + def test_extracts_score_when_think_start_tag_missing(self) -> None: + """Test extraction when model omits and starts with .""" + logprobs = [ + _make_logprob(""), + _make_logprob("\n"), + _make_logprob("<"), + _make_logprob("score"), + _make_logprob(">"), + _make_logprob(" no", top_logprobs=[]), + _make_logprob(" "), + ] + result = _extract_tokens_inside_score_tag(logprobs) + assert result.token == " no" + + +class TestGetRiskyProbabilities: + """Tests for _get_risky_probabilities.""" + + def test_returns_risky_probability(self) -> None: + """Test normalized risky probability computation.""" + token = _make_token_logprob( + "no", + top_logprobs=[ + ("no", math.log(0.8)), + ("yes", math.log(0.2)), + ], + ) + result = _get_risky_probabilities(token) + assert result == pytest.approx(0.2, abs=1e-6) + + def test_all_safe(self) -> None: + """Test when all probability goes to safe token.""" + token = _make_token_logprob( + "no", + top_logprobs=[ + ("no", math.log(0.99)), + ("No", math.log(0.01)), + ], + ) + result = _get_risky_probabilities(token) + assert result == pytest.approx(0.0, abs=1e-6) + + def test_all_risky(self) -> None: + """Test when all probability goes to risky token.""" + token = _make_token_logprob( + "yes", + top_logprobs=[ + ("yes", math.log(0.95)), + ("Yes", math.log(0.05)), + ], + ) + result = _get_risky_probabilities(token) + assert result == pytest.approx(1.0, abs=1e-6) + + def test_raises_on_underflow(self) -> None: + """Test that zero total probability raises UnexpectedModelBehavior.""" + token = _make_token_logprob( + "maybe", + top_logprobs=[ + ("maybe", -1.0), + ("perhaps", -2.0), + ], + ) + with pytest.raises(UnexpectedModelBehavior, match="underflow"): + _get_risky_probabilities(token) + + def test_handles_whitespace_in_tokens(self) -> None: + """Test that tokens with surrounding whitespace are normalized.""" + token = _make_token_logprob( + " no ", + top_logprobs=[ + (" no ", math.log(0.7)), + (" yes ", math.log(0.3)), + ], + ) + result = _get_risky_probabilities(token) + assert result == pytest.approx(0.3, abs=1e-6) + + def test_handles_case_insensitive_tokens(self) -> None: + """Test that token matching is case-insensitive.""" + token = _make_token_logprob( + "No", + top_logprobs=[ + ("No", math.log(0.6)), + ("YES", math.log(0.4)), + ], + ) + result = _get_risky_probabilities(token) + assert result == pytest.approx(0.4, abs=1e-6) + + +class TestIsSafe: + """Tests for is_safe.""" + + def _make_full_logprobs(self, safe_prob: float, risky_prob: float): + """Build a complete logprobs list for a single-token score response.""" + return [ + _make_logprob(""), + _make_logprob("\n"), + _make_logprob(""), + _make_logprob(""), + _make_logprob( + "no", + logprob=math.log(safe_prob), + top_logprobs=[ + {"token": "no", "logprob": math.log(safe_prob), "bytes": None}, + {"token": "yes", "logprob": math.log(risky_prob), "bytes": None}, + ], + ), + _make_logprob(""), + ] + + def test_safe_when_below_threshold(self) -> None: + """Test that input is safe when risky probability is below threshold.""" + logprobs = self._make_full_logprobs(0.9, 0.1) + assert is_safe(0.5, logprobs) is True + + def test_unsafe_when_above_threshold(self) -> None: + """Test that input is unsafe when risky probability exceeds threshold.""" + logprobs = self._make_full_logprobs(0.3, 0.7) + assert is_safe(0.5, logprobs) is False + + def test_unsafe_when_at_threshold(self) -> None: + """Test that input is unsafe when risky probability equals threshold.""" + logprobs = self._make_full_logprobs(0.5, 0.5) + assert is_safe(0.5, logprobs) is False diff --git a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py index 2b7a25930..f35a0a551 100644 --- a/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py +++ b/tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capability.py @@ -22,80 +22,82 @@ SUBJECT_ALLOWED, SUBJECT_REJECTED, QuestionValidity, - _extract_conversation_id, - _extract_message_str_from_user_content, +) +from pydantic_ai_lightspeed.capabilities.utils import ( + extract_conversation_id, + extract_message_str_from_user_content, ) _MODULE = "pydantic_ai_lightspeed.capabilities.question_validity._capability" class TestExtractMessageStrFromUserContent: - """Tests for _extract_message_str_from_user_content helper.""" + """Tests for extract_message_str_from_user_content helper.""" def test_extracts_plain_strings(self) -> None: """Test extraction from a sequence of plain strings.""" content = ["hello", "world"] - result = _extract_message_str_from_user_content(content) + result = extract_message_str_from_user_content(content) assert result == "hello\nworld" def test_extracts_text_content(self) -> None: """Test extraction from TextContent objects.""" content = [TextContent(content="first"), TextContent(content="second")] - result = _extract_message_str_from_user_content(content) + result = extract_message_str_from_user_content(content) assert result == "first\nsecond" def test_mixed_str_and_text_content(self) -> None: """Test extraction from a mix of strings and TextContent.""" content = ["plain", TextContent(content="rich")] - result = _extract_message_str_from_user_content(content) + result = extract_message_str_from_user_content(content) assert result == "plain\nrich" def test_empty_sequence(self) -> None: """Test extraction from an empty sequence.""" - result = _extract_message_str_from_user_content([]) + result = extract_message_str_from_user_content([]) assert result == "" def test_single_string(self) -> None: """Test extraction from a single-element sequence.""" - result = _extract_message_str_from_user_content(["only"]) + result = extract_message_str_from_user_content(["only"]) assert result == "only" def test_sequence_with_non_text_content(self) -> None: """Test extraction from a single-element sequence.""" - result = _extract_message_str_from_user_content([ImageUrl("fake.png"), "keep"]) + result = extract_message_str_from_user_content([ImageUrl("fake.png"), "keep"]) assert result == "keep" class TestExtractConversationId: - """Tests for _extract_conversation_id helper.""" + """Tests for extract_conversation_id helper.""" def test_extracts_conversation_id(self, mocker: MockerFixture) -> None: """Test extraction when extra_body.conversation is set.""" model = mocker.Mock() model.settings = {"extra_body": {"conversation": "conv_123"}} - assert _extract_conversation_id(model) == "conv_123" + assert extract_conversation_id(model) == "conv_123" def test_returns_none_when_settings_missing(self, mocker: MockerFixture) -> None: """Test that None settings yields None.""" model = mocker.Mock() model.settings = None - assert _extract_conversation_id(model) is None + assert extract_conversation_id(model) is None def test_returns_none_when_extra_body_missing(self, mocker: MockerFixture) -> None: """Test that missing extra_body yields None.""" model = mocker.Mock() model.settings = {} - assert _extract_conversation_id(model) is None + assert extract_conversation_id(model) is None def test_returns_none_when_extra_body_not_dict(self, mocker: MockerFixture) -> None: """Test that a non-dict extra_body yields None instead of raising.""" model = mocker.Mock() model.settings = {"extra_body": "not-a-dict"} - assert _extract_conversation_id(model) is None + assert extract_conversation_id(model) is None def test_returns_none_when_conversation_not_string( self, mocker: MockerFixture @@ -104,7 +106,7 @@ def test_returns_none_when_conversation_not_string( model = mocker.Mock() model.settings = {"extra_body": {"conversation": 123}} - assert _extract_conversation_id(model) is None + assert extract_conversation_id(model) is None class TestQuestionValidityConfigInit: