review comment

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

Author: akihikokuroda

AGENTS.md

@@ -180,8 +180,8 @@ Intrinsics are specialized LoRA adapters that add task-specific capabilities (RA
 | `rag` | `find_citations(response, documents, context, backend)` | Document sentences supporting the response |
 | `rag` | `check_context_relevance(question, document, context, backend)` | Whether a document is relevant (0–1); only supported for granite-4.0, not granite-4.1 |
 | `rag` | `flag_hallucinated_content(response, documents, context, backend)` | Flag potentially hallucinated sentences |
-| `guardian` | `factuality_detection(response, context, backend, *, documents=None, model_options=None)` | Determine if the last response is factually incorrect ("yes"/"no") |
-| `guardian` | `factuality_correction(response, context, backend, *, documents=None, model_options=None)` | Correct the last response to be factually accurate |
+| `guardian` | `factuality_detection(context, backend, *, documents=None, model_options=None)` | Determine if the last response is factually incorrect ("yes"/"no") |
+| `guardian` | `factuality_correction(context, backend, *, documents=None, model_options=None)` | Correct the last response to be factually accurate |
 
 ```python
 from mellea.backends.huggingface import LocalHFBackend

mellea/stdlib/components/intrinsic/guardian.py

@@ -179,7 +179,15 @@ def guardian_check(
 
     Returns:
         Risk score as a float between 0.0 (no risk) and 1.0 (risk detected).
+
+    Raises:
+        ValueError: If documents are provided but target_role is not "assistant".
     """
+    if documents is not None and target_role != "assistant":
+        raise ValueError(
+            "documents parameter is only supported when target_role='assistant'"
+        )
+
     if documents is not None and target_role == "assistant":
         response, context, resolved_docs = _resolve_response(None, context)
         explicit_docs = _coerce_to_documents(documents)

test/stdlib/components/intrinsic/test_guardian.py

@@ -182,12 +182,24 @@ def test_factuality_detection(backend):
 
 @pytest.mark.qualitative
 def test_factuality_detection_from_context(backend):
-    """Verify factuality detection works when documents are in the last message."""
+    """Verify factuality detection works when documents are already in the last message."""
     context, documents = _read_guardian_input("factuality_detection.json")
 
-    # Test with documents passed as argument to factuality_detection
-    # The function should handle extracting the response and attaching documents
-    result = guardian.factuality_detection(context, backend, documents=documents)
+    # Rebuild context with documents attached to the last assistant message
+    last_turn = context.last_turn()
+    if last_turn and last_turn.model_input:
+        context = ChatContext().add(Message("user", "What is the question?"))
+        if isinstance(last_turn.model_input, Message):
+            context = context.add(
+                Message(
+                    last_turn.model_input.role,
+                    last_turn.model_input.content,
+                    documents=documents,
+                )
+            )
+
+    # Call without documents= argument; documents should be picked up from context
+    result = guardian.factuality_detection(context, backend)
     assert result == "yes" or result == "no"
 
 
@@ -203,12 +215,24 @@ def test_factuality_correction(backend):
 
 @pytest.mark.qualitative
 def test_factuality_correction_from_context(backend):
-    """Verify factuality correction works when documents are in the last message."""
+    """Verify factuality correction works when documents are already in the last message."""
     context, documents = _read_guardian_input("factuality_correction.json")
 
-    # Test with documents passed as argument to factuality_correction
-    # The function should handle extracting the response and attaching documents
-    result = guardian.factuality_correction(context, backend, documents=documents)
+    # Rebuild context with documents attached to the last assistant message
+    last_turn = context.last_turn()
+    if last_turn and last_turn.model_input:
+        context = ChatContext().add(Message("user", "What is the question?"))
+        if isinstance(last_turn.model_input, Message):
+            context = context.add(
+                Message(
+                    last_turn.model_input.role,
+                    last_turn.model_input.content,
+                    documents=documents,
+                )
+            )
+
+    # Call without documents= argument; documents should be picked up from context
+    result = guardian.factuality_correction(context, backend)
     assert isinstance(result, str)