LCORE-1830: Implement Question Validity Safety Capability in Pydantic AI#1913
LCORE-1830: Implement Question Validity Safety Capability in Pydantic AI#1913Jazzcort wants to merge 1 commit into
Conversation
WalkthroughAdds QuestionValidity, an AbstractCapability that builds a classification prompt from ctx.prompt, calls a configured model to allow/reject the prompt, strips conversation settings from the model, updates token usage, and short-circuits the agent run on rejection. ChangesQuestion Validity Capability
🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/pydantic_ai_lightspeed/capabilities/__init__.py`:
- Line 8: Replace the relative re-exports with absolute imports: in
src/pydantic_ai_lightspeed/capabilities/__init__.py change the relative import
of QuestionValidity to the package-absolute path (importing QuestionValidity
from pydantic_ai_lightspeed.capabilities.question_validity), and apply the same
pattern in src/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py
for any internal re-exports; update the import statements that currently use the
“from .module import Name” form to “from
pydantic_ai_lightspeed.capabilities.module import Name” so they follow the
project’s absolute-import guideline.
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py`:
- Line 13: Replace usage of the stdlib logger with the project logger factory:
remove the "import logging" and instead import get_logger from the project's log
helper and create a module-level logger via logger = get_logger(__name__)
(replace any logging.getLogger(...) occurrences in this module, e.g., where the
current logger is referenced at the top-level and around the symbol that
initializes the logger on line ~25). Ensure all subsequent logger calls use this
module-level logger variable.
- Line 71: Add Google-style docstrings and complete type annotations for all
functions in this module: specifically document
_extract_message_str_from_user_content with a short description, Args, Returns
and Raises sections and any other module functions referenced near the top and
middle of the file; explicitly annotate __post_init__ with -> None and add a
Google-style docstring describing its post-initialization behavior; ensure every
function/method has full parameter and return type hints and descriptive
docstrings (follow Google Python docstring conventions).
- Line 152: The guardrail comparison currently does an exact match (result.text
== SUBJECT_ALLOWED) which fails on casing/whitespace variants; change the check
to normalize the classifier output first (e.g., normalized =
result.text.strip().upper()) and compare normalized against SUBJECT_ALLOWED and
SUBJECT_REJECTED (or compare against normalized versions of those constants)
inside the same function/method where result.text is used; also add unit tests
in TestWrapRun exercising variants like "allowed", " ALLOWED\n", and mixed-case
to ensure normalization works.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 5d0b2f52-fac2-4f56-8e38-bc92319852c6
📒 Files selected for processing (6)
src/pydantic_ai_lightspeed/capabilities/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.pytests/unit/pydantic_ai_lightspeed/capabilities/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)
- GitHub Check: E2E Tests for Lightspeed Evaluation job
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E: server mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 3
🧰 Additional context used
📓 Path-based instructions (3)
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Usepytest.mark.asynciomarker for async tests
Files:
tests/unit/pydantic_ai_lightspeed/capabilities/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Llama Stack imports: Usefrom llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Useasync deffor I/O operations and external API calls
Use standard log levels with clear purposes:debug()for diagnostic info,info()for program execution,warning()for unexpected events,error()for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes:Configuration,Error/Exception,Resolver,Interface
Abstract classes must use ABC with@abstractmethoddecorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes
Files:
src/pydantic_ai_lightspeed/capabilities/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py
src/**/__init__.py
📄 CodeRabbit inference engine (AGENTS.md)
Package
__init__.pyfiles must contain brief package descriptions
Files:
src/pydantic_ai_lightspeed/capabilities/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py
🔇 Additional comments (3)
tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py (1)
1-520: LGTM!tests/unit/pydantic_ai_lightspeed/capabilities/__init__.py (1)
1-2: LGTM!tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py (1)
1-2: LGTM!
| SUBJECT_ALLOWED = "ALLOWED" | ||
|
|
||
|
|
||
| def _extract_message_str_from_user_content(user_content: Sequence[UserContent]) -> str: |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win
Add missing function docstrings and complete annotations.
Multiple functions in this module are missing descriptive docstrings, and __post_init__ (Line 123) is missing an explicit -> None return annotation.
As per coding guidelines, src/**/*.py: “All functions must have complete type annotations for parameters and return types … and include descriptive docstrings” and “Follow Google Python docstring conventions with required sections”.
Also applies to: 83-83, 123-123, 126-126, 139-141
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py` at
line 71, Add Google-style docstrings and complete type annotations for all
functions in this module: specifically document
_extract_message_str_from_user_content with a short description, Args, Returns
and Raises sections and any other module functions referenced near the top and
middle of the file; explicitly annotate __post_init__ with -> None and add a
Google-style docstring describing its post-initialization behavior; ensure every
function/method has full parameter and return type hints and descriptive
docstrings (follow Google Python docstring conventions).
Source: Coding guidelines
| # Include token usage from the question validity request | ||
| ctx.usage.incr(result.usage) | ||
|
|
||
| if result.text == SUBJECT_ALLOWED: |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify guardrail tests currently cover normalization behavior.
rg -n -C2 'SUBJECT_ALLOWED|result\.text|wrap_run' tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 4763
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact comparison at the referenced line range
python - <<'PY'
from pathlib import Path
path = Path("src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py")
lines = path.read_text().splitlines()
start = 130
end = 175
for i in range(start, min(end, len(lines))):
print(f"{i+1:4d}: {lines[i]}")
PY
# Check tests for normalization/whitespace/casing expectations
rg -n "SUBJECT_ALLOWED|SUBJECT_REJECTED|strip\(|upper\(|lower\(|whitespace|casefold|normalize" \
tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 2123
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# 1) Inspect the wrap_run tests area
python - <<'PY'
from pathlib import Path
path = Path("tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py")
lines = path.read_text().splitlines()
# Print a focused window around the wrap_run tests
start = 240
end = 560
for i in range(start, min(end, len(lines))):
print(f"{i+1:4d}: {lines[i]}")
PY
# 2) Search for any normalization on result.text in this capability
rg -n "result\.text|strip\(|upper\(|lower\(|casefold\(|normalize" \
src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py src/pydantic_ai_lightspeed/capabilities/question_validity -S
# 3) Search the tests for any variants (e.g., whitespace/case) around SUBJECT_ALLOWED/REJECTED
rg -n "TextPart\(content=\s*['\"]\s*(SUBJECT_ALLOWED|SUBJECT_REJECTED)" tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py
rg -n "SUBJECT_ALLOWED\s*(\)|\]|,)|SUBJECT_REJECTED\s*(\)|\]|,)" tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 12090
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate model_request implementation (within this capability if present)
rg -n "def model_request|async def model_request|model_request\s*\(" -S src/pydantic_ai_lightspeed/capabilities/question_validity
# Show the top of _capacity.py to understand how result.text is produced/used
python - <<'PY'
from pathlib import Path
path = Path("src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py")
lines = path.read_text().splitlines()
for i in range(1, min(220, len(lines))+1):
if i<=40 or (140<=i<=180):
print(f"{i:4d}: {lines[i-1]}")
PY
# Search for other SUBJECT_ALLOWED comparisons
rg -n "result\.text\s*==\s*SUBJECT_ALLOWED|SUBJECT_ALLOWED\s*in\s*result\.text|SUBJECT_ALLOWED" \
src/pydantic_ai_lightspeed/capabilities/question_validity -S
# Search tests for any normalization or casing/whitespace cases
rg -n "SUBJECT_ALLOWED.*(strip|upper|lower|casefold)|SUBJECT_REJECTED.*(strip|upper|lower|casefold)" \
tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 3233
Normalize guardrail classifier output before ALLOWED/REJECTED comparison.
src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py line 152 does an exact result.text == SUBJECT_ALLOWED match, so casing/whitespace variants (e.g., "allowed", " ALLOWED\n") would be treated as rejected. Current TestWrapRun unit tests only exercise exact SUBJECT_ALLOWED/SUBJECT_REJECTED strings (and a non-matching string), with no coverage for whitespace/case normalization.
Suggested fix
- if result.text == SUBJECT_ALLOWED:
+ classification = result.text.strip().upper()
+ if classification == SUBJECT_ALLOWED:
return await handler() # proceed with the real run
else:
# short-circuit: return the rejection message with shield usage tracked
state = GraphAgentState(usage=ctx.usage)
return AgentRunResult(output=self.invalid_question_response, _state=state)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py` at
line 152, The guardrail comparison currently does an exact match (result.text ==
SUBJECT_ALLOWED) which fails on casing/whitespace variants; change the check to
normalize the classifier output first (e.g., normalized =
result.text.strip().upper()) and compare normalized against SUBJECT_ALLOWED and
SUBJECT_REJECTED (or compare against normalized versions of those constants)
inside the same function/method where result.text is used; also add unit tests
in TestWrapRun exercising variants like "allowed", " ALLOWED\n", and mixed-case
to ensure normalization works.
Jazzcort
left a comment
There was a problem hiding this comment.
The comment is just some bugs that I found when I tried to wire up this question validity shield with pydantic agent. 😁
| return "\n".join(str_arr) | ||
|
|
||
|
|
||
| def _remove_conversation_from_settings(model: Model) -> Model: |
There was a problem hiding this comment.
During my testing, if we share Model either through passing the same Model instance as what is passed to Agent in build_agent or grabbing it from ctx.model, the request made by question validity shield contaminate the stored conversation history in Llama Stack. Therefore, the solutions that I can think of are either we copy the Model instance with "conversation" field removed from "extra_body" like this or we pass a separate Model instance without "conversation" field in "extra_body" when we create this capability in build_agent.
fe2e532 to
3700804
Compare
Implement an LLM-based guardrail that classifies user questions as on-topic (Kubernetes/OpenShift or customized topic) before the main agent processes them. Off-topic questions are short-circuited with a rejection message, bypassing the primary agent entirely. Includes unit tests.
3700804 to
9599ca5
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (2)
src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py (2)
158-158:⚠️ Potential issue | 🟠 Major | ⚡ Quick winNormalize the classifier output before the ALLOWED check.
Line 158 does an exact sentinel match, so harmless casing or whitespace differences from the model will incorrectly short-circuit valid questions as rejected.
[suggested fix]
Minimal fix
- if result.text == SUBJECT_ALLOWED: + classification = result.text.strip().upper() + if classification == SUBJECT_ALLOWED: return await handler() # proceed with the real run else: # short-circuit: return the rejection message with shield usage tracked state = GraphAgentState(usage=ctx.usage) return AgentRunResult(output=self.invalid_question_response, _state=state)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py` at line 158, Normalize the classifier output before comparing to the sentinel: ensure result.text is non-null, strip surrounding whitespace and case-normalize it (e.g., lower()) and compare that normalized string to a similarly normalized SUBJECT_ALLOWED (or normalize SUBJECT_ALLOWED once). Update the check around result.text == SUBJECT_ALLOWED in the capacity logic (the result.text comparison in _capacity.py) to use the normalized value so harmless casing/whitespace differences won’t short-circuit valid questions.
72-89: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Google-style docstrings to the helpers and capability methods, and update the class description.
_extract_message_str_from_user_content,_remove_conversation_from_settings,__post_init__,_build_prompt, andwrap_rundo not meet the repository’s required docstring format, and theQuestionValidityclass docstring still describes a generic boolean-returning guard instead of this capability’s actual short-circuit behavior.As per coding guidelines,
src/**/*.py: “All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings” and “Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes”.Also applies to: 105-118, 129-147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py` around lines 72 - 89, Update the docstrings and type annotations to follow repository standards: add Google-style docstrings (with Parameters, Returns, Raises as appropriate) to the helper functions _extract_message_str_from_user_content(user_content: Sequence[UserContent]) -> str and _remove_conversation_from_settings(model: Model) -> Model, and to the methods __post_init__(self) -> None, _build_prompt(self, ...) -> str, and wrap_run(self, ...) -> Any (use PEP 604 union syntax where needed, e.g., str | None), ensuring all parameter and return types are fully annotated; also revise the QuestionValidity class docstring to accurately describe its short-circuit behavior (not a generic boolean guard) and include an Attributes section for key fields. Ensure docstrings succinctly document parameters, return values, raised exceptions, and any side effects without changing logic.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py`:
- Line 158: Normalize the classifier output before comparing to the sentinel:
ensure result.text is non-null, strip surrounding whitespace and case-normalize
it (e.g., lower()) and compare that normalized string to a similarly normalized
SUBJECT_ALLOWED (or normalize SUBJECT_ALLOWED once). Update the check around
result.text == SUBJECT_ALLOWED in the capacity logic (the result.text comparison
in _capacity.py) to use the normalized value so harmless casing/whitespace
differences won’t short-circuit valid questions.
- Around line 72-89: Update the docstrings and type annotations to follow
repository standards: add Google-style docstrings (with Parameters, Returns,
Raises as appropriate) to the helper functions
_extract_message_str_from_user_content(user_content: Sequence[UserContent]) ->
str and _remove_conversation_from_settings(model: Model) -> Model, and to the
methods __post_init__(self) -> None, _build_prompt(self, ...) -> str, and
wrap_run(self, ...) -> Any (use PEP 604 union syntax where needed, e.g., str |
None), ensuring all parameter and return types are fully annotated; also revise
the QuestionValidity class docstring to accurately describe its short-circuit
behavior (not a generic boolean guard) and include an Attributes section for key
fields. Ensure docstrings succinctly document parameters, return values, raised
exceptions, and any side effects without changing logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0c6cff2a-26dd-4d38-bff2-bdac045dfd76
📒 Files selected for processing (6)
src/pydantic_ai_lightspeed/capabilities/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.pytests/unit/pydantic_ai_lightspeed/capabilities/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: server mode / ci / group 3
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E Tests for Lightspeed Evaluation job
🧰 Additional context used
📓 Path-based instructions (3)
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Usepytest.mark.asynciomarker for async tests
Files:
tests/unit/pydantic_ai_lightspeed/capabilities/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Llama Stack imports: Usefrom llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Useasync deffor I/O operations and external API calls
Use standard log levels with clear purposes:debug()for diagnostic info,info()for program execution,warning()for unexpected events,error()for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes:Configuration,Error/Exception,Resolver,Interface
Abstract classes must use ABC with@abstractmethoddecorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes
Files:
src/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pysrc/pydantic_ai_lightspeed/capabilities/__init__.pysrc/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py
src/**/__init__.py
📄 CodeRabbit inference engine (AGENTS.md)
Package
__init__.pyfiles must contain brief package descriptions
Files:
src/pydantic_ai_lightspeed/capabilities/question_validity/__init__.pysrc/pydantic_ai_lightspeed/capabilities/__init__.py
🔇 Additional comments (6)
tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py (1)
1-520: LGTM!tests/unit/pydantic_ai_lightspeed/capabilities/__init__.py (1)
1-1: LGTM!tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py (1)
1-1: LGTM!src/pydantic_ai_lightspeed/capabilities/__init__.py (1)
1-10: LGTM!src/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py (1)
1-7: LGTM!src/pydantic_ai_lightspeed/capabilities/question_validity/_capacity.py (1)
1-69: LGTM!
Description
Implement an LLM-based guardrail that classifies user questions as on-topic (Kubernetes/OpenShift or customized topic) before the main agent processes them. Off-topic questions are short-circuited with a rejection message, bypassing the primary agent entirely. Includes unit tests.
More thorough tests (e2e) can be implemented when we wire up the pydantic AI agent to our endpoints and see if this question validity shield really shields the agent from off-topic questions.
Type of change
Tools used to create PR
The unit tests were drafted by Claude Code which was then reviewed and refined by me. 😁
Related Tickets & Documents
LCORE-1830
Checklist before requesting a review
Testing
Since we have not wired pydantic_ai to our endpoints, you can only do an ad-hoc test with it.
Summary by CodeRabbit
New Features
Tests