review comments

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

Author: akihikokuroda

docs/examples/intrinsics/factuality_correction.py

@@ -86,7 +86,9 @@
 )
 # NOTE: This example can also be run with the OpenAIBackend using a GraniteSwitch model. See docs/examples/granite-switch/.
 
-ctx = ctx.add(Message("user", user_text)).add(Message("assistant", response_text))
+ctx = ctx.add(Message("user", user_text))
 
-result = guardian.factuality_correction(ctx, backend, documents=[document])
+result = guardian.factuality_correction(
+    response_text, ctx, backend, documents=[document]
+)
 print(f"Result of factuality correction: {result}")  # corrected response string

docs/examples/intrinsics/factuality_detection.py

@@ -29,7 +29,9 @@
 )
 # NOTE: This example can also be run with the OpenAIBackend using a GraniteSwitch model. See docs/examples/granite-switch/.
 
-ctx = ctx.add(Message("user", user_text)).add(Message("assistant", response_text))
+ctx = ctx.add(Message("user", user_text))
 
-result = guardian.factuality_detection(ctx, backend, documents=[document])
+result = guardian.factuality_detection(
+    response_text, ctx, backend, documents=[document]
+)
 print(f"Result of factuality detection: {result}")  # string "yes" or "no"

mellea/stdlib/components/intrinsic/_util.py

@@ -13,6 +13,7 @@
 from ....stdlib import functional as mfuncs
 from ...components import Document
 from ...context import ChatContext
+from ..chat import Message
 from .intrinsic import Intrinsic
 
 
@@ -74,20 +75,37 @@ def _resolve_response(
     """Return ``(response_text, context_to_use)``.
 
     When *response* is not ``None``, returns it with *context* unchanged.
-    When ``None``, extracts from the last turn's ``output.value`` and rewinds
-    *context* to before that output.
+    When ``None``, extracts from the last turn's ``output.value`` (generated) or
+    ``model_input.content`` (manually-added Message), then rewinds *context*
+    to before that turn.
     """
     if response is not None:
         return response, context
     turn = context.last_turn()
-    if turn is None or turn.output is None:
+    if turn is None:
         raise ValueError("response is None and context has no last turn with output")
-    if turn.output.value is None:
-        raise ValueError("response is None and last turn output has no value")
+
+    # Try generated output first
+    if turn.output is not None:
+        if turn.output.value is None:
+            raise ValueError("response is None and last turn output has no value")
+        response_text = turn.output.value
+    # Fall back to manually-added assistant Message
+    elif (
+        turn.model_input is not None
+        and isinstance(turn.model_input, Message)
+        and turn.model_input.role == "assistant"
+    ):
+        response_text = turn.model_input.content
+    else:
+        raise ValueError(
+            "response is None and context has no last turn with output or assistant message"
+        )
+
     rewound = context.previous_node
     if rewound is None:
         raise ValueError("Cannot rewind context past the root node")
-    return turn.output.value, rewound  # type: ignore[return-value]
+    return response_text, rewound  # type: ignore[return-value]
 
 
 def call_intrinsic(

mellea/stdlib/components/intrinsic/guardian.py

@@ -7,15 +7,14 @@
 """
 
 import collections.abc
-from typing import cast
 
 from ....backends import model_ids
 from ....backends.adapters import AdapterMixin
 from ...components import Document
 from ...context import ChatContext
 from ..chat import Message
 from ..docs.document import _coerce_to_documents
-from ._util import call_intrinsic
+from ._util import _resolve_response, call_intrinsic
 
 
 def policy_guardrails(
@@ -153,6 +152,8 @@ def guardian_check(
     backend: AdapterMixin,
     criteria: str,
     target_role: str = "assistant",
+    *,
+    documents: collections.abc.Iterable[str | Document] | None = None,
     model_options: dict | None = None,
 ) -> float:
     """Check whether text meets specified safety/quality criteria.
@@ -168,12 +169,20 @@ def guardian_check(
             criteria string.
         target_role: Role whose last message is being evaluated
             (``"user"`` or ``"assistant"``).
+        documents: Optional document snippets to attach to the target message.
+            Primarily used for the ``"groundedness"`` criterion, to provide
+            reference context for grounding checks. Each element may be a
+            ``Document`` or a plain string (automatically wrapped in ``Document``).
+            Keyword-only.
         model_options: Optional model options to pass to the backend (e.g.,
             temperature, max_tokens). Defaults to ``{ModelOption.TEMPERATURE: 0.0}``.
 
     Returns:
         Risk score as a float between 0.0 (no risk) and 1.0 (risk detected).
     """
+    if documents is not None and target_role == "assistant":
+        context = _reattach_documents(context, documents)
+
     criteria_text = CRITERIA_BANK.get(criteria, criteria)
 
     scoring = (
@@ -209,61 +218,57 @@ def _reattach_documents(
         New context with documents attached to the last assistant message.
 
     Raises:
-        ValueError: If context cannot be rewound or assistant content cannot be extracted.
+        ValueError: If context cannot be rewound or content cannot be extracted.
     """
-    last_turn = context.last_turn()
-    if last_turn is None:
-        raise ValueError("Cannot reattach documents: context has no last turn")
-
-    # Extract assistant content, preferring generated output over input
-    if last_turn.output is not None and last_turn.output.value is not None:
-        assistant_content = last_turn.output.value
-    elif last_turn.output is not None and last_turn.output.value is None:
-        # Uncomputed thunk — avoid silent fallthrough to model_input
-        raise ValueError(
-            "Cannot reattach documents: last turn output is uncomputed (thunk with no value)"
-        )
-    elif last_turn.model_input is not None and isinstance(
-        last_turn.model_input, Message
+    turn = context.last_turn()
+    if turn is None:
+        raise ValueError("Cannot extract response from empty context")
+
+    # Try to get response from output first (generated), then from message content (manual)
+    response_text = None
+    rewound = context.previous_node
+
+    if turn.output is not None and turn.output.value is not None:
+        # Response is from a generated output
+        response_text = turn.output.value
+    elif (
+        turn.model_input is not None
+        and isinstance(turn.model_input, Message)
+        and turn.model_input.role == "assistant"
     ):
-        assistant_content = last_turn.model_input.content
+        # Response is from a manually added assistant Message
+        response_text = turn.model_input.content
     else:
         raise ValueError(
-            "Cannot reattach documents: cannot extract assistant content from last turn"
+            "Cannot extract response: turn has neither output nor assistant message"
         )
 
-    # Rewind and re-add with documents
-    rewound = context.previous_node
     if rewound is None:
         raise ValueError("Cannot rewind context past the root node")
 
-    return cast(
-        ChatContext,
-        rewound.add(
-            Message(
-                "assistant",
-                assistant_content,
-                documents=_coerce_to_documents(documents),
-            )
-        ),
+    return rewound.add(  # type: ignore[return-value]
+        Message("assistant", response_text, documents=_coerce_to_documents(documents))
     )
 
 
 def factuality_detection(
+    response: str | None,
     context: ChatContext,
     backend: AdapterMixin,
     *,
     documents: collections.abc.Iterable[str | Document] | None = None,
     model_options: dict | None = None,
 ) -> str:
-    """Determine if the last response is factually incorrect.
+    """Determine if a response is factually incorrect.
 
-    Intrinsic function that evaluates the factuality of the
-    assistant's response to a user's question. The context should end with
-    a user question followed by an assistant answer.
+    Intrinsic function that evaluates the factuality of an assistant's response
+    to a user's question. The context should typically end with a user question
+    followed by an assistant answer.
 
     Args:
-        context: Chat context containing user question and assistant answer.
+        response: The assistant's response text to evaluate. When ``None``, the
+            response is extracted from the last assistant output in ``context``.
+        context: Chat context containing user question and conversation history.
         backend: Backend instance that supports LoRA/aLoRA adapters.
         documents: Document snippets that provide factual context for evaluation.
             Each element may be a ``Document`` or a plain string (automatically
@@ -283,8 +288,24 @@ def factuality_detection(
 ### Scoring Schema: If the last assistant's text meets the criteria, return 'yes'; otherwise, return 'no'.
 """
 
+    if response is None:
+        response, context = _resolve_response(None, context)
+
     if documents is not None:
-        context = _reattach_documents(context, documents)
+        if response is not None:
+            # Response was explicitly provided, add it with documents
+            context = context.add(
+                Message(
+                    "assistant", response, documents=_coerce_to_documents(documents)
+                )
+            )
+        else:
+            # Response came from context output, reattach documents
+            context = _reattach_documents(context, documents)
+    else:
+        # No documents provided, add response to context if it was explicitly provided
+        if response is not None:
+            context = context.add(Message("assistant", response))
 
     context = context.add(Message("user", detector_message))
     result_json = call_intrinsic(
@@ -294,19 +315,22 @@ def factuality_detection(
 
 
 def factuality_correction(
+    response: str | None,
     context: ChatContext,
     backend: AdapterMixin,
     *,
     documents: collections.abc.Iterable[str | Document] | None = None,
     model_options: dict | None = None,
 ) -> str:
-    """Corrects the last response so that it is factually correct.
+    """Correct a response to be factually accurate.
 
     Intrinsic function that corrects the assistant's response to a user's
     question relative to the given contextual information.
 
     Args:
-        context: Chat context containing user question and assistant answer.
+        response: The assistant's response text to correct. When ``None``, the
+            response is extracted from the last assistant output in ``context``.
+        context: Chat context containing user question and conversation history.
         backend: Backend instance that supports LoRA/aLoRA adapters.
         documents: Document snippets that provide factual context for correction.
             Each element may be a ``Document`` or a plain string (automatically
@@ -326,8 +350,24 @@ def factuality_correction(
 ### Scoring Schema: If the last assistant's text meets the criteria, return a corrected version of the assistant's message based on the given context; otherwise, return 'none'.
 """
 
+    if response is None:
+        response, context = _resolve_response(None, context)
+
     if documents is not None:
-        context = _reattach_documents(context, documents)
+        if response is not None:
+            # Response was explicitly provided, add it with documents
+            context = context.add(
+                Message(
+                    "assistant", response, documents=_coerce_to_documents(documents)
+                )
+            )
+        else:
+            # Response came from context output, reattach documents
+            context = _reattach_documents(context, documents)
+    else:
+        # No documents provided, add response to context if it was explicitly provided
+        if response is not None:
+            context = context.add(Message("assistant", response))
 
     context = context.add(Message("user", corrector_message))
     result_json = call_intrinsic(

test/backends/test_openai_intrinsics.py

@@ -559,7 +559,9 @@ def test_call_intrinsic_factuality_detection(call_intrinsic_backend):
     for m in messages:
         context = context.add(Message(m["role"], m["content"]))
 
-    result = guardian.factuality_detection(docs, context, call_intrinsic_backend)
+    result = guardian.factuality_detection(
+        None, context, call_intrinsic_backend, documents=docs
+    )
     assert result in ("yes", "no")
 
 
@@ -580,5 +582,7 @@ def test_call_intrinsic_factuality_correction(call_intrinsic_backend):
     for m in messages:
         context = context.add(Message(m["role"], m["content"]))
 
-    result = guardian.factuality_correction(docs, context, call_intrinsic_backend)
+    result = guardian.factuality_correction(
+        None, context, call_intrinsic_backend, documents=docs
+    )
     assert isinstance(result, str)

test/stdlib/components/intrinsic/test_guardian.py

@@ -176,7 +176,7 @@ def test_factuality_detection(backend):
     context, documents = _read_guardian_input("factuality_detection.json")
 
     # Test with documents passed as argument
-    result = guardian.factuality_detection(context, backend, documents=documents)
+    result = guardian.factuality_detection(None, context, backend, documents=documents)
     assert result == "yes" or result == "no"
 
 
@@ -185,23 +185,11 @@ def test_factuality_detection_from_context(backend):
     """Verify factuality detection works when documents are in the last message."""
     context, documents = _read_guardian_input("factuality_detection.json")
 
-    # Extract assistant content using the same logic as _reattach_documents
-    last_turn = context.last_turn()
-    assert last_turn is not None
-    if last_turn.output is not None and last_turn.output.value is not None:
-        content = last_turn.output.value
-    else:
-        assert isinstance(last_turn.model_input, Message)
-        content = last_turn.model_input.content
-
-    rewound = context.previous_node
-    assert rewound is not None
-    context_with_docs: ChatContext = rewound.add(  # type: ignore[assignment]
-        Message("assistant", content, documents=documents)
-    )
+    # Extract response and rewind, then add back with documents
+    context_with_docs = guardian._reattach_documents(context, documents)
 
-    # Test with documents=None (should extract from context)
-    result = guardian.factuality_detection(context_with_docs, backend)
+    # Test with response=None (should extract from context)
+    result = guardian.factuality_detection(None, context_with_docs, backend)
     assert result == "yes" or result == "no"
 
 
@@ -211,7 +199,7 @@ def test_factuality_correction(backend):
     context, documents = _read_guardian_input("factuality_correction.json")
 
     # Test with documents passed as argument
-    result = guardian.factuality_correction(context, backend, documents=documents)
+    result = guardian.factuality_correction(None, context, backend, documents=documents)
     assert isinstance(result, str)
 
 
@@ -220,23 +208,11 @@ def test_factuality_correction_from_context(backend):
     """Verify factuality correction works when documents are in the last message."""
     context, documents = _read_guardian_input("factuality_correction.json")
 
-    # Extract assistant content using the same logic as _reattach_documents
-    last_turn = context.last_turn()
-    assert last_turn is not None
-    if last_turn.output is not None and last_turn.output.value is not None:
-        content = last_turn.output.value
-    else:
-        assert isinstance(last_turn.model_input, Message)
-        content = last_turn.model_input.content
-
-    rewound = context.previous_node
-    assert rewound is not None
-    context_with_docs: ChatContext = rewound.add(  # type: ignore[assignment]
-        Message("assistant", content, documents=documents)
-    )
+    # Extract response and rewind, then add back with documents
+    context_with_docs = guardian._reattach_documents(context, documents)
 
-    # Test with documents=None (should extract from context)
-    result = guardian.factuality_correction(context_with_docs, backend)
+    # Test with response=None (should extract from context)
+    result = guardian.factuality_correction(None, context_with_docs, backend)
     assert isinstance(result, str)