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
WalkthroughThis PR adds ChangesQuestion Validity Capability
Sequence Diagram(s)sequenceDiagram
participant RunContext
participant QuestionValidity
participant model_request
participant handler
participant AgentRunResult
RunContext->>QuestionValidity: wrap_run(ctx, handler=...)
QuestionValidity->>QuestionValidity: _build_prompt(ctx.prompt)
QuestionValidity->>model_request: classify user input (built prompt)
model_request-->>QuestionValidity: SUBJECT_ALLOWED or SUBJECT_REJECTED
alt SUBJECT_ALLOWED
QuestionValidity->>handler: invoke provided handler
handler-->>AgentRunResult: return result
else SUBJECT_REJECTED
QuestionValidity->>AgentRunResult: short-circuit with invalid_question_response
end
QuestionValidity->>RunContext: update usage with classification tokens
AgentRunResult-->>RunContext: return result
🎯 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
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.
Jazzcort
force-pushed
the
question-validity-for-pydantic-ai
branch
2 times, most recently
from
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!
jrobertboos
left a comment
jrobertboos
left a comment
There was a problem hiding this comment.
Instead of naming the files capacity they should be capability.
|
@jrobertboos My bad! Thanks for catching that! 😁 |
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.
Jazzcort
force-pushed
the
question-validity-for-pydantic-ai
branch
from
9599ca5 to
7e60d1c
Compare
|
Changed! @jrobertboos Thanks for reviewing! |
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (7)
src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py (7)
129-130: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd docstring to
__post_init__method.Project guidelines require all functions to have descriptive docstrings.
🤖 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/_capability.py` around lines 129 - 130, The __post_init__ method lacks a docstring; add a concise docstring above def __post_init__(self) describing its purpose (e.g., that it sanitizes/normalizes the model settings by removing conversation-related keys), mention that it mutates self.model and calls _remove_conversation_from_settings(self.model), and note expected behavior/return type (None) and any side effects so it complies with project docstring guidelines.Source: Coding guidelines
72-82: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Google-style docstring.
The function lacks the required Args, Returns, and Raises sections specified in project guidelines.
🤖 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/_capability.py` around lines 72 - 82, The function _extract_message_str_from_user_content is missing a Google-style docstring; update its docstring to include a short summary plus Args (describe user_content: Sequence[UserContent]), Returns (str: combined text joined by newlines), and Raises (document that no exceptions are raised or note specific exceptions if any can propagate), following the project's Google-style format; reference UserContent and TextContent in the description so readers know which variants are handled and keep the content concise and accurate.Source: Coding guidelines
145-163: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Google-style docstring to
wrap_runmethod.Project guidelines require all functions to have descriptive docstrings with Args, Returns, and Raises sections.
🤖 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/_capability.py` around lines 145 - 163, Add a Google-style docstring to the async method wrap_run that describes its purpose and behaviour (it builds a prompt, calls model_request, increments ctx.usage with the result usage, and either calls the provided handler to continue or short-circuits with an invalid response), and include Args (ctx: RunContext, handler: WrapRunHandler), Returns (AgentRunResult) and Raises (any exceptions propagated from model_request or handler) sections; reference the side-effect of ctx.usage.incr(result.usage) and mention the returned GraphAgentState when short-circuiting.Source: Coding guidelines
132-143: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Google-style docstring to
_build_promptmethod.Project guidelines require all functions to have descriptive docstrings with Args, Returns, and Raises sections.
🤖 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/_capability.py` around lines 132 - 143, Add a Google-style docstring to the _build_prompt method describing its purpose, parameters, return value, and potential errors: explain the message parameter (type: str | Sequence[UserContent] | None) and that it may be converted via _extract_message_str_from_user_content, state that it returns a str (the prompt created by substituting message, allowed=SUBJECT_ALLOWED, rejected=SUBJECT_REJECTED into self.model_prompt using Template), and document any Raises (e.g., Template-related errors or TypeError if message contains unsupported types). Place this docstring immediately above the def _build_prompt declaration and reference Template, self.model_prompt, and _extract_message_str_from_user_content in the description.Source: Coding guidelines
158-158:⚠️ Potential issue | 🟠 Major | ⚡ Quick winNormalize classifier output before comparison.
The exact match
result.text == SUBJECT_ALLOWEDwill fail if the LLM returns whitespace variants (e.g.," ALLOWED\n") or different casing (e.g.,"allowed"), causing valid inputs to be incorrectly rejected.🤖 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/_capability.py` at line 158, The equality check using `if result.text == SUBJECT_ALLOWED` should normalize the classifier output to avoid failing on whitespace or case differences; update the condition in `_capability.py` (where `result.text` is compared to `SUBJECT_ALLOWED`) to compare a normalized form such as `result.text.strip().casefold()` against the normalized `SUBJECT_ALLOWED` (e.g., `SUBJECT_ALLOWED.casefold()`) so trailing/leading whitespace and casing differences no longer cause false negatives.
85-101: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Google-style docstring.
The function lacks the required Args, Returns, and Raises sections specified in project guidelines.
🤖 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/_capability.py` around lines 85 - 101, The function _remove_conversation_from_settings lacks a Google-style docstring with Args, Returns, and Raises sections; add a docstring above the function following Google style that documents: Args (model: Model) — describe that a Model is accepted and that its settings may be mutated/copied; Returns — a Model instance (either the original or a shallow-copied Model with "conversation" removed from settings["extra_body"]); Raises — any exceptions that could be propagated (e.g., TypeError if model.settings is not a mapping) or state that no exceptions are raised under normal operation. Reference the function name _remove_conversation_from_settings and the Model type in the docstring so readers can locate and understand the behavior.Source: Coding guidelines
106-118: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winAdd Attributes section to class docstring.
Project guidelines require class docstrings to include an Attributes section describing the class fields.
🤖 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/_capability.py` around lines 106 - 118, Update the class docstring for QuestionValidity to include an Attributes section that lists each class field (name and type), a one-line description of its purpose, and default value if any; locate the docstring for the QuestionValidity class in src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py and add entries for all attributes defined on the class (e.g., model, any guard or config fields, etc.), following the project's docstring style and formatting.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.
Inline comments:
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py`:
- Line 99: The code is directly mutating the private attribute _model._settings;
instead, avoid touching private internals—use the Model's public API to produce
a modified instance (e.g., a copy/with/replace_settings method on the Model) and
set the new settings there, or if such an API is unavailable, create a shallow
clone of _model with the updated settings via its public constructor or copy
method and use that; if you must mutate the private field as a last resort, add
a clear comment documenting why and the risk and wrap the mutation in a single
well-scoped helper function (referencing _model and _settings to locate the
change).
- Line 32: Update the prompt string in _capability.py that currently reads "Your
job is to determine where or a user's question is related to kubernetes and/or
openshift technologies and to provide a one-word response." to fix the typo by
replacing "where or" with "whether" so it reads "Your job is to determine
whether a user's question is related to kubernetes and/or openshift technologies
and to provide a one-word response."; locate the prompt constant or variable
containing that text in the Question Validity capability and perform the
substitution.
In
`@tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py`:
- Around line 67-70: Update the test docstring for
test_sequence_with_non_text_content to accurately describe its behavior: state
that it verifies non-TextContent items (e.g., ImageUrl("fake.png")) are filtered
out and text items are preserved when calling
_extract_message_str_from_user_content, rather than claiming it tests a
"single-element sequence"; reference the test function name
test_sequence_with_non_text_content and the helper
_extract_message_str_from_user_content in the updated docstring for clarity.
---
Duplicate comments:
In `@src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py`:
- Around line 129-130: The __post_init__ method lacks a docstring; add a concise
docstring above def __post_init__(self) describing its purpose (e.g., that it
sanitizes/normalizes the model settings by removing conversation-related keys),
mention that it mutates self.model and calls
_remove_conversation_from_settings(self.model), and note expected
behavior/return type (None) and any side effects so it complies with project
docstring guidelines.
- Around line 72-82: The function _extract_message_str_from_user_content is
missing a Google-style docstring; update its docstring to include a short
summary plus Args (describe user_content: Sequence[UserContent]), Returns (str:
combined text joined by newlines), and Raises (document that no exceptions are
raised or note specific exceptions if any can propagate), following the
project's Google-style format; reference UserContent and TextContent in the
description so readers know which variants are handled and keep the content
concise and accurate.
- Around line 145-163: Add a Google-style docstring to the async method wrap_run
that describes its purpose and behaviour (it builds a prompt, calls
model_request, increments ctx.usage with the result usage, and either calls the
provided handler to continue or short-circuits with an invalid response), and
include Args (ctx: RunContext, handler: WrapRunHandler), Returns
(AgentRunResult) and Raises (any exceptions propagated from model_request or
handler) sections; reference the side-effect of ctx.usage.incr(result.usage) and
mention the returned GraphAgentState when short-circuiting.
- Around line 132-143: Add a Google-style docstring to the _build_prompt method
describing its purpose, parameters, return value, and potential errors: explain
the message parameter (type: str | Sequence[UserContent] | None) and that it may
be converted via _extract_message_str_from_user_content, state that it returns a
str (the prompt created by substituting message, allowed=SUBJECT_ALLOWED,
rejected=SUBJECT_REJECTED into self.model_prompt using Template), and document
any Raises (e.g., Template-related errors or TypeError if message contains
unsupported types). Place this docstring immediately above the def _build_prompt
declaration and reference Template, self.model_prompt, and
_extract_message_str_from_user_content in the description.
- Line 158: The equality check using `if result.text == SUBJECT_ALLOWED` should
normalize the classifier output to avoid failing on whitespace or case
differences; update the condition in `_capability.py` (where `result.text` is
compared to `SUBJECT_ALLOWED`) to compare a normalized form such as
`result.text.strip().casefold()` against the normalized `SUBJECT_ALLOWED` (e.g.,
`SUBJECT_ALLOWED.casefold()`) so trailing/leading whitespace and casing
differences no longer cause false negatives.
- Around line 85-101: The function _remove_conversation_from_settings lacks a
Google-style docstring with Args, Returns, and Raises sections; add a docstring
above the function following Google style that documents: Args (model: Model) —
describe that a Model is accepted and that its settings may be mutated/copied;
Returns — a Model instance (either the original or a shallow-copied Model with
"conversation" removed from settings["extra_body"]); Raises — any exceptions
that could be propagated (e.g., TypeError if model.settings is not a mapping) or
state that no exceptions are raised under normal operation. Reference the
function name _remove_conversation_from_settings and the Model type in the
docstring so readers can locate and understand the behavior.
- Around line 106-118: Update the class docstring for QuestionValidity to
include an Attributes section that lists each class field (name and type), a
one-line description of its purpose, and default value if any; locate the
docstring for the QuestionValidity class in
src/pydantic_ai_lightspeed/capabilities/question_validity/_capability.py and add
entries for all attributes defined on the class (e.g., model, any guard or
config fields, etc.), following the project's docstring style and formatting.
🪄 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: 9931376f-ec7f-4429-9dd3-3bab0bcdca00
📒 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/_capability.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: server mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E: server mode / ci / group 2
🧰 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/question_validity/__init__.pytests/unit/pydantic_ai_lightspeed/capabilities/__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/_capability.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 (11)
src/pydantic_ai_lightspeed/capabilities/__init__.py (1)
1-10: LGTM!src/pydantic_ai_lightspeed/capabilities/question_validity/__init__.py (1)
1-7: 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!tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py (7)
1-22: LGTM!
24-34: LGTM!
39-66: LGTM!
73-143: LGTM!
145-189: LGTM!
191-242: LGTM!
245-519: LGTM!
| Instructions: | ||
| - You are a question classifying tool | ||
| - You are an expert in kubernetes and openshift | ||
| - Your job is to determine where or a user's question is related to kubernetes and/or openshift technologies and to provide a one-word response. |
There was a problem hiding this comment.
Fix typo in classification prompt.
Line 32 contains "where or a user's question" which should be "whether a user's question".
📝 Suggested fix
-- Your job is to determine where or a user's question is related to kubernetes and/or openshift technologies and to provide a one-word response.
+- Your job is to determine whether a user's question is related to kubernetes and/or openshift technologies and to provide a one-word response.🤖 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/_capability.py` at
line 32, Update the prompt string in _capability.py that currently reads "Your
job is to determine where or a user's question is related to kubernetes and/or
openshift technologies and to provide a one-word response." to fix the typo by
replacing "where or" with "whether" so it reads "Your job is to determine
whether a user's question is related to kubernetes and/or openshift technologies
and to provide a one-word response."; locate the prompt constant or variable
containing that text in the Question Validity capability and perform the
substitution.
| _settings = copy(settings) | ||
| _settings["extra_body"] = _extra_body | ||
| _model = copy(model) | ||
| _model._settings = _settings |
There was a problem hiding this comment.
Avoid mutating private attribute _settings.
Directly assigning to _model._settings couples the code to Pydantic AI's internal implementation. If the Model class changes how it stores settings, this will break.
💡 Consider safer alternatives
If Model provides a public API for copying with modified settings, use that. Otherwise, document the risk and add a comment explaining why direct mutation is necessary.
🤖 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/_capability.py` at
line 99, The code is directly mutating the private attribute _model._settings;
instead, avoid touching private internals—use the Model's public API to produce
a modified instance (e.g., a copy/with/replace_settings method on the Model) and
set the new settings there, or if such an API is unavailable, create a shallow
clone of _model with the updated settings via its public constructor or copy
method and use that; if you must mutate the private field as a last resort, add
a clear comment documenting why and the risk and wrap the mutation in a single
well-scoped helper function (referencing _model and _settings to locate the
change).
| 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"]) | ||
| assert result == "keep" |
There was a problem hiding this comment.
Fix the test docstring to match the test behavior.
The docstring says "Test extraction from a single-element sequence." but the test actually validates that non-TextContent items (like ImageUrl) are filtered out while text content is preserved. This appears to be a copy-paste error from the previous test.
📝 Suggested docstring fix
- def test_sequence_with_non_text_content(self) -> None:
- """Test extraction from a single-element sequence."""
+ def test_sequence_with_non_text_content(self) -> None:
+ """Test that non-TextContent items are skipped during extraction."""
result = _extract_message_str_from_user_content([ImageUrl("fake.png"), "keep"])
assert result == "keep"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| 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"]) | |
| assert result == "keep" | |
| def test_sequence_with_non_text_content(self) -> None: | |
| """Test that non-TextContent items are skipped during extraction.""" | |
| result = _extract_message_str_from_user_content([ImageUrl("fake.png"), "keep"]) | |
| assert result == "keep" |
🤖 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
`@tests/unit/pydantic_ai_lightspeed/capabilities/question_validity/test_capacity.py`
around lines 67 - 70, Update the test docstring for
test_sequence_with_non_text_content to accurately describe its behavior: state
that it verifies non-TextContent items (e.g., ImageUrl("fake.png")) are filtered
out and text items are preserved when calling
_extract_message_str_from_user_content, rather than claiming it tests a
"single-element sequence"; reference the test function name
test_sequence_with_non_text_content and the helper
_extract_message_str_from_user_content in the updated docstring for clarity.
2 participants
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