feat: fail early for citations with incompatible models

Author: theagenticguy

src/strands/models/bedrock.py

@@ -18,7 +18,11 @@
 from ..event_loop import streaming
 from ..tools import convert_pydantic_to_tool_spec
 from ..types.content import ContentBlock, Message, Messages
-from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
+from ..types.exceptions import (
+    ContextWindowOverflowException,
+    ModelThrottledException,
+    UnsupportedModelCitationsException,
+)
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolResult, ToolSpec
 from .model import Model
@@ -34,6 +38,15 @@
     "too many total text bytes",
 ]
 
+# Model IDs that support citation functionality
+CITATION_SUPPORTED_MODELS = [
+    "anthropic.claude-3-5-sonnet-20241022-v2:0",
+    "anthropic.claude-3-7-sonnet-20250219-v1:0",
+    "anthropic.claude-opus-4-20250514-v1:0",
+    "anthropic.claude-sonnet-4-20250514-v1:0",
+    "anthropic.claude-opus-4-1-20250805-v1:0",
+]
+
 T = TypeVar("T", bound=BaseModel)
 
 
@@ -349,6 +362,39 @@ def _generate_redaction_events(self) -> list[StreamEvent]:
 
         return events
 
+    def _has_citations_config(self, messages: Messages) -> bool:
+        """Check if any message contains document content with citations enabled.
+
+        Args:
+            messages: List of messages to check for citations config.
+
+        Returns:
+            True if any message contains a document with citations enabled, False otherwise.
+        """
+        for message in messages:
+            for content_block in message["content"]:
+                if "document" in content_block:
+                    document = content_block["document"]
+                    if "citations" in document and document["citations"] is not None:
+                        citations_config = document["citations"]
+                        if "enabled" in citations_config and citations_config["enabled"]:
+                            return True
+        return False
+
+    def _validate_citations_support(self, messages: Messages) -> None:
+        """Validate that the current model supports citations if citations are requested.
+
+        Args:
+            messages: List of messages to check for citations config.
+
+        Raises:
+            UnsupportedModelCitationsException: If citations are requested but the model doesn't support them.
+        """
+        if self._has_citations_config(messages):
+            model_id = self.config["model_id"]
+            if model_id not in CITATION_SUPPORTED_MODELS:
+                raise UnsupportedModelCitationsException(model_id, CITATION_SUPPORTED_MODELS)
+
     @override
     async def stream(
         self,
@@ -374,7 +420,10 @@ async def stream(
         Raises:
             ContextWindowOverflowException: If the input exceeds the model's context window.
             ModelThrottledException: If the model service is throttling requests.
+            UnsupportedModelCitationsException: If citations are requested but the model doesn't support them.
         """
+        # Validate citations support before starting the thread (fail fast in async context)
+        self._validate_citations_support(messages)
 
         def callback(event: Optional[StreamEvent] = None) -> None:
             loop.call_soon_threadsafe(queue.put_nowait, event)

src/strands/types/exceptions.py

@@ -75,3 +75,27 @@ class SessionException(Exception):
     """Exception raised when session operations fail."""
 
     pass
+
+
+class UnsupportedModelCitationsException(Exception):
+    """Exception raised when trying to use citations with an unsupported model.
+
+    This exception is raised when a user attempts to use document citations with a Bedrock model
+    that does not support the citations feature. Citations are only supported by specific Claude models.
+    """
+
+    def __init__(self, model_id: str, supported_models: list[str]) -> None:
+        """Initialize exception with model information.
+
+        Args:
+            model_id: The model ID that doesn't support citations.
+            supported_models: List of model IDs that do support citations.
+        """
+        self.model_id = model_id
+        self.supported_models = supported_models
+
+        message = (
+            f"Model '{model_id}' does not support document citations. "
+            f"Supported models for citations are: {', '.join(supported_models)}"
+        )
+        super().__init__(message)

tests/strands/models/test_bedrock.py

@@ -12,7 +12,7 @@
 import strands
 from strands.models import BedrockModel
 from strands.models.bedrock import DEFAULT_BEDROCK_MODEL_ID, DEFAULT_BEDROCK_REGION
-from strands.types.exceptions import ModelThrottledException
+from strands.types.exceptions import ModelThrottledException, UnsupportedModelCitationsException
 from strands.types.tools import ToolSpec
 
 
@@ -1228,3 +1228,205 @@ def test_format_request_cleans_tool_result_content_blocks(model, model_id):
     assert tool_result == expected
     assert "extraField" not in tool_result
     assert "mcpMetadata" not in tool_result
+
+
+# Citation validation tests
+
+
+def test_has_citations_config_with_enabled_citations(bedrock_client):
+    """Test _has_citations_config returns True when citations are enabled."""
+    model = BedrockModel()
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": True},
+                    }
+                }
+            ],
+        }
+    ]
+
+    assert model._has_citations_config(messages) is True
+
+
+def test_has_citations_config_with_disabled_citations(bedrock_client):
+    """Test _has_citations_config returns False when citations are disabled."""
+    model = BedrockModel()
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": False},
+                    }
+                }
+            ],
+        }
+    ]
+
+    assert model._has_citations_config(messages) is False
+
+
+def test_has_citations_config_with_no_citations(bedrock_client):
+    """Test _has_citations_config returns False when no citations config."""
+    model = BedrockModel()
+    messages = [
+        {
+            "role": "user",
+            "content": [{"document": {"name": "test.pdf", "source": {"bytes": b"test content"}, "format": "pdf"}}],
+        }
+    ]
+
+    assert model._has_citations_config(messages) is False
+
+
+def test_has_citations_config_with_no_documents(bedrock_client):
+    """Test _has_citations_config returns False when no documents."""
+    model = BedrockModel()
+    messages = [{"role": "user", "content": [{"text": "test message"}]}]
+
+    assert model._has_citations_config(messages) is False
+
+
+def test_validate_citations_support_with_supported_model(bedrock_client):
+    """Test _validate_citations_support passes with supported model."""
+    model = BedrockModel(model_id="anthropic.claude-3-5-sonnet-20241022-v2:0")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": True},
+                    }
+                }
+            ],
+        }
+    ]
+
+    # Should not raise an exception
+    model._validate_citations_support(messages)
+
+
+def test_validate_citations_support_with_unsupported_model(bedrock_client):
+    """Test _validate_citations_support raises exception with unsupported model."""
+    model = BedrockModel(model_id="us.anthropic.claude-3-5-haiku-20241022-v1:0")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": True},
+                    }
+                }
+            ],
+        }
+    ]
+
+    with pytest.raises(UnsupportedModelCitationsException) as exc_info:
+        model._validate_citations_support(messages)
+
+    assert "us.anthropic.claude-3-5-haiku-20241022-v1:0" in str(exc_info.value)
+    assert "anthropic.claude-3-5-sonnet-20241022-v2:0" in str(exc_info.value)
+
+
+def test_validate_citations_support_with_no_citations(bedrock_client):
+    """Test _validate_citations_support passes when no citations are requested."""
+    model = BedrockModel(model_id="us.anthropic.claude-3-5-haiku-20241022-v1:0")
+    messages = [{"role": "user", "content": [{"text": "test message"}]}]
+
+    # Should not raise an exception since no citations are requested
+    model._validate_citations_support(messages)
+
+
+@pytest.mark.asyncio
+async def test_stream_with_citations_unsupported_model(bedrock_client, alist):
+    """Test that stream raises exception for unsupported model with citations."""
+    model = BedrockModel(model_id="us.anthropic.claude-3-5-haiku-20241022-v1:0")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": True},
+                    }
+                },
+                {"text": "What does this document say?"},
+            ],
+        }
+    ]
+
+    with pytest.raises(UnsupportedModelCitationsException) as exc_info:
+        await alist(model.stream(messages))
+
+    assert "us.anthropic.claude-3-5-haiku-20241022-v1:0" in str(exc_info.value)
+    assert "anthropic.claude-3-5-sonnet-20241022-v2:0" in str(exc_info.value)
+
+    # Verify the Bedrock client was not called since validation failed early
+    bedrock_client.converse_stream.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_stream_with_citations_supported_model(bedrock_client, alist):
+    """Test that stream works with supported model and citations."""
+    model = BedrockModel(model_id="anthropic.claude-opus-4-20250514-v1:0")
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "document": {
+                        "name": "test.pdf",
+                        "source": {"bytes": b"test content"},
+                        "format": "pdf",
+                        "citations": {"enabled": True},
+                    }
+                },
+                {"text": "What does this document say?"},
+            ],
+        }
+    ]
+
+    # Mock successful response
+    bedrock_client.converse_stream.return_value = {"stream": [{"messageStart": {"role": "assistant"}}]}
+
+    # Should not raise an exception
+    await alist(model.stream(messages))
+
+    # Verify the Bedrock client was called since validation passed
+    bedrock_client.converse_stream.assert_called_once()
+
+
+def test_unsupported_model_citations_exception_message():
+    """Test that UnsupportedModelCitationsException has proper message format."""
+    model_id = "us.anthropic.claude-3-5-haiku-20241022-v1:0"
+    supported_models = ["anthropic.claude-3-5-sonnet-20241022-v2:0"]
+
+    exception = UnsupportedModelCitationsException(model_id, supported_models)
+
+    assert model_id in str(exception)
+    assert "anthropic.claude-3-5-sonnet-20241022-v2:0" in str(exception)
+    assert "does not support document citations" in str(exception)
+    assert exception.model_id == model_id
+    assert exception.supported_models == supported_models