diff --git a/lib/crewai/src/crewai/llms/providers/bedrock/__init__.py b/lib/crewai/src/crewai/llms/providers/bedrock/__init__.py index e69de29bb..57d2d43e1 100644 --- a/lib/crewai/src/crewai/llms/providers/bedrock/__init__.py +++ b/lib/crewai/src/crewai/llms/providers/bedrock/__init__.py @@ -0,0 +1,12 @@ +from crewai.llms.providers.bedrock.completion import ( + BedrockCitation, + BedrockCompletion, + BedrockGroundedResponse, +) + + +__all__ = [ + "BedrockCitation", + "BedrockCompletion", + "BedrockGroundedResponse", +] diff --git a/lib/crewai/src/crewai/llms/providers/bedrock/completion.py b/lib/crewai/src/crewai/llms/providers/bedrock/completion.py index 47946d949..e22aabbbd 100644 --- a/lib/crewai/src/crewai/llms/providers/bedrock/completion.py +++ b/lib/crewai/src/crewai/llms/providers/bedrock/completion.py @@ -30,7 +30,6 @@ if TYPE_CHECKING: SystemContentBlockTypeDef, TokenUsageTypeDef, ToolConfigurationTypeDef, - ToolTypeDef, ) from crewai.llms.hooks.base import BaseInterceptor @@ -49,6 +48,38 @@ except ImportError: STRUCTURED_OUTPUT_TOOL_NAME = "structured_output" +class BedrockCitation(TypedDict, total=False): + """Type definition for a citation from Bedrock Web Grounding. + + Citations are returned when using system tools like nova_grounding + that provide web-grounded responses with source attribution. + """ + + url: str + domain: str + + +class BedrockGroundedResponse(TypedDict): + """Type definition for a grounded response with citations. + + When system tools like nova_grounding are used, responses include + both the generated text and citations to source material. + """ + + text: str + citations: list[BedrockCitation] + + +class SystemToolTypeDef(TypedDict): + """Type definition for a Bedrock system tool. + + System tools are built-in tools provided by AWS Bedrock that the model + can use internally, such as nova_grounding for web search with citations. + """ + + name: str + + def _preprocess_structured_data( data: dict[str, Any], response_model: type[BaseModel] ) -> dict[str, Any]: @@ -246,6 +277,7 @@ class BedrockCompletion(BaseLLM): additional_model_response_field_paths: list[str] | None = None, interceptor: BaseInterceptor[Any, Any] | None = None, response_format: type[BaseModel] | None = None, + system_tools: list[str] | None = None, **kwargs: Any, ) -> None: """Initialize AWS Bedrock completion client. @@ -268,6 +300,10 @@ class BedrockCompletion(BaseLLM): interceptor: HTTP interceptor (not yet supported for Bedrock). response_format: Pydantic model for structured output. Used as default when response_model is not passed to call()/acall() methods. + system_tools: List of Bedrock system tool names to enable (e.g., ["nova_grounding"]). + System tools are built-in tools that the model can use internally. + When nova_grounding is enabled, responses may include citations + from web sources. See AWS documentation for available system tools. **kwargs: Additional parameters """ if interceptor is not None: @@ -332,6 +368,7 @@ class BedrockCompletion(BaseLLM): self.additional_model_response_field_paths = ( additional_model_response_field_paths ) + self.system_tools = system_tools # Model-specific settings self.is_claude_model = "claude" in model.lower() @@ -412,26 +449,29 @@ class BedrockCompletion(BaseLLM): cast(object, [{"text": system_message}]), ) - # Add tool config if present or if messages contain tool content - # Bedrock requires toolConfig when messages have toolUse/toolResult + # Build tool config with regular tools and system tools + tool_config_tools: list[Any] = [] + if tools: - tool_config: ToolConfigurationTypeDef = { - "tools": cast( - "Sequence[ToolTypeDef]", - cast(object, self._format_tools_for_converse(tools)), - ) - } - body["toolConfig"] = tool_config + tool_config_tools.extend(self._format_tools_for_converse(tools)) elif self._messages_contain_tool_content(formatted_messages): - # Create minimal toolConfig from tool history in messages tools_from_history = self._extract_tools_from_message_history( formatted_messages ) if tools_from_history: - body["toolConfig"] = cast( - "ToolConfigurationTypeDef", - cast(object, {"tools": tools_from_history}), - ) + tool_config_tools.extend(tools_from_history) + + # Add system tools (e.g., nova_grounding) if configured + if self.system_tools: + tool_config_tools.extend( + {"systemTool": {"name": name}} for name in self.system_tools + ) + + if tool_config_tools: + body["toolConfig"] = cast( + "ToolConfigurationTypeDef", + cast(object, {"tools": tool_config_tools}), + ) # Add optional advanced features if configured if self.guardrail_config: @@ -543,26 +583,31 @@ class BedrockCompletion(BaseLLM): cast(object, [{"text": system_message}]), ) - # Add tool config if present or if messages contain tool content - # Bedrock requires toolConfig when messages have toolUse/toolResult + # Build tool config with regular tools and system tools + async_tool_config_tools: list[Any] = [] + if tools: - tool_config: ToolConfigurationTypeDef = { - "tools": cast( - "Sequence[ToolTypeDef]", - cast(object, self._format_tools_for_converse(tools)), - ) - } - body["toolConfig"] = tool_config + async_tool_config_tools.extend( + self._format_tools_for_converse(tools) + ) elif self._messages_contain_tool_content(formatted_messages): - # Create minimal toolConfig from tool history in messages tools_from_history = self._extract_tools_from_message_history( formatted_messages ) if tools_from_history: - body["toolConfig"] = cast( - "ToolConfigurationTypeDef", - cast(object, {"tools": tools_from_history}), - ) + async_tool_config_tools.extend(tools_from_history) + + # Add system tools (e.g., nova_grounding) if configured + if self.system_tools: + async_tool_config_tools.extend( + {"systemTool": {"name": name}} for name in self.system_tools + ) + + if async_tool_config_tools: + body["toolConfig"] = cast( + "ToolConfigurationTypeDef", + cast(object, {"tools": async_tool_config_tools}), + ) if self.guardrail_config: guardrail_config: GuardrailConfigurationTypeDef = cast( @@ -766,12 +811,29 @@ class BedrockCompletion(BaseLLM): # Process content blocks and handle tool use correctly text_content = "" + all_citations: list[BedrockCitation] = [] for content_block in content: # Handle text content if "text" in content_block: text_content += content_block["text"] + # Extract citations from citationsContent if present + # (returned by system tools like nova_grounding) + if "citationsContent" in content_block: + citations_content = content_block["citationsContent"] + citations_list = citations_content.get("citations", []) + for citation in citations_list: + location = citation.get("location", {}) + web_info = location.get("web", {}) + if web_info: + all_citations.append( + BedrockCitation( + url=web_info.get("url", ""), + domain=web_info.get("domain", ""), + ) + ) + # Handle tool use - corrected structure according to AWS API docs elif "toolUse" in content_block and available_functions: tool_use_block = content_block["toolUse"] @@ -842,12 +904,23 @@ class BedrockCompletion(BaseLLM): messages=messages, ) - return self._invoke_after_llm_call_hooks( + # Apply after hooks to the text content + final_text = self._invoke_after_llm_call_hooks( messages, text_content, from_agent, ) + # If citations were extracted (from system tools like nova_grounding), + # return a BedrockGroundedResponse with both text and citations + if all_citations: + return BedrockGroundedResponse( + text=final_text, + citations=all_citations, + ) + + return final_text + except ClientError as e: # Handle all AWS ClientError exceptions as per documentation error_code = e.response.get("Error", {}).get("Code", "Unknown") @@ -1352,11 +1425,28 @@ class BedrockCompletion(BaseLLM): return non_structured_output_tool_uses text_content = "" + async_all_citations: list[BedrockCitation] = [] for content_block in content: if "text" in content_block: text_content += content_block["text"] + # Extract citations from citationsContent if present + # (returned by system tools like nova_grounding) + if "citationsContent" in content_block: + citations_content = content_block["citationsContent"] + citations_list = citations_content.get("citations", []) + for citation in citations_list: + location = citation.get("location", {}) + web_info = location.get("web", {}) + if web_info: + async_all_citations.append( + BedrockCitation( + url=web_info.get("url", ""), + domain=web_info.get("domain", ""), + ) + ) + elif "toolUse" in content_block and available_functions: tool_use_block = content_block["toolUse"] tool_use_id = tool_use_block.get("toolUseId") @@ -1424,6 +1514,14 @@ class BedrockCompletion(BaseLLM): messages=messages, ) + # If citations were extracted (from system tools like nova_grounding), + # return a BedrockGroundedResponse with both text and citations + if async_all_citations: + return BedrockGroundedResponse( + text=text_content, + citations=async_all_citations, + ) + return text_content except ClientError as e: diff --git a/lib/crewai/tests/llms/bedrock/test_bedrock.py b/lib/crewai/tests/llms/bedrock/test_bedrock.py index efe3191e7..830daa471 100644 --- a/lib/crewai/tests/llms/bedrock/test_bedrock.py +++ b/lib/crewai/tests/llms/bedrock/test_bedrock.py @@ -948,3 +948,282 @@ def test_bedrock_agent_kickoff_structured_output_with_tools(): assert result.pydantic.result == 42, f"Expected result 42 but got {result.pydantic.result}" assert result.pydantic.operation, "Operation should not be empty" assert result.pydantic.explanation, "Explanation should not be empty" + + +# ============================================================================= +# System Tools (nova_grounding) Tests +# ============================================================================= + + +def test_bedrock_system_tools_initialization(): + """ + Test that BedrockCompletion can be initialized with system_tools parameter. + """ + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + from crewai.llms.providers.bedrock.completion import BedrockCompletion + assert isinstance(llm, BedrockCompletion) + assert llm.system_tools == ["nova_grounding"] + + +def test_bedrock_system_tools_in_tool_config(bedrock_mocks): + """ + Test that system tools are properly included in the toolConfig when calling the API. + """ + _, mock_client = bedrock_mocks + + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + # Call the LLM + llm.call("What is the latest news about AI?") + + # Verify the API was called with system tools in toolConfig + mock_client.converse.assert_called_once() + call_kwargs = mock_client.converse.call_args[1] + + assert "toolConfig" in call_kwargs + tool_config = call_kwargs["toolConfig"] + assert "tools" in tool_config + + # Check that system tool is included + tools = tool_config["tools"] + system_tool_found = False + for tool in tools: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "System tool nova_grounding not found in toolConfig" + + +def test_bedrock_citation_extraction(bedrock_mocks): + """ + Test that citations are properly extracted from Bedrock responses with citationsContent. + """ + _, mock_client = bedrock_mocks + + # Mock response with citations (as returned by nova_grounding) + mock_response_with_citations = { + 'output': { + 'message': { + 'role': 'assistant', + 'content': [ + { + 'text': 'According to recent reports, AI has made significant advances.', + 'citationsContent': { + 'citations': [ + { + 'location': { + 'web': { + 'url': 'https://example.com/ai-news', + 'domain': 'example.com' + } + } + }, + { + 'location': { + 'web': { + 'url': 'https://techsite.com/article', + 'domain': 'techsite.com' + } + } + } + ] + } + } + ] + } + }, + 'usage': { + 'inputTokens': 10, + 'outputTokens': 20, + 'totalTokens': 30 + } + } + mock_client.converse.return_value = mock_response_with_citations + + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + result = llm.call("What is the latest news about AI?") + + # Result should be a BedrockGroundedResponse with citations + from crewai.llms.providers.bedrock.completion import BedrockGroundedResponse + assert isinstance(result, dict), f"Expected dict (BedrockGroundedResponse), got {type(result)}" + assert "text" in result + assert "citations" in result + assert result["text"] == "According to recent reports, AI has made significant advances." + assert len(result["citations"]) == 2 + assert result["citations"][0]["url"] == "https://example.com/ai-news" + assert result["citations"][0]["domain"] == "example.com" + assert result["citations"][1]["url"] == "https://techsite.com/article" + assert result["citations"][1]["domain"] == "techsite.com" + + +def test_bedrock_no_citations_returns_string(bedrock_mocks): + """ + Test that responses without citations return a plain string (backward compatibility). + """ + _, mock_client = bedrock_mocks + + # Standard response without citations + mock_response_no_citations = { + 'output': { + 'message': { + 'role': 'assistant', + 'content': [ + {'text': 'This is a regular response without citations.'} + ] + } + }, + 'usage': { + 'inputTokens': 10, + 'outputTokens': 5, + 'totalTokens': 15 + } + } + mock_client.converse.return_value = mock_response_no_citations + + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + result = llm.call("Hello, how are you?") + + # Result should be a plain string when no citations are present + assert isinstance(result, str), f"Expected str, got {type(result)}" + assert result == "This is a regular response without citations." + + +def test_bedrock_system_tools_combined_with_regular_tools(bedrock_mocks): + """ + Test that system tools can be combined with regular tools in the same request. + """ + _, mock_client = bedrock_mocks + + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + # Define a regular tool + regular_tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "The city name"} + }, + "required": ["location"] + } + } + } + ] + + llm.call("What's the weather in New York?", tools=regular_tools) + + # Verify both regular tools and system tools are in toolConfig + mock_client.converse.assert_called_once() + call_kwargs = mock_client.converse.call_args[1] + + assert "toolConfig" in call_kwargs + tools = call_kwargs["toolConfig"]["tools"] + + # Should have both regular tool and system tool + has_regular_tool = any("toolSpec" in t for t in tools) + has_system_tool = any("systemTool" in t for t in tools) + + assert has_regular_tool, "Regular tool not found in toolConfig" + assert has_system_tool, "System tool not found in toolConfig" + + +def test_bedrock_grounded_response_types_exported(): + """ + Test that BedrockCitation and BedrockGroundedResponse types are properly exported. + """ + from crewai.llms.providers.bedrock import ( + BedrockCitation, + BedrockCompletion, + BedrockGroundedResponse, + ) + + # Verify the types are accessible + assert BedrockCitation is not None + assert BedrockGroundedResponse is not None + assert BedrockCompletion is not None + + # Verify they are TypedDicts with expected keys + assert "url" in BedrockCitation.__annotations__ + assert "domain" in BedrockCitation.__annotations__ + assert "text" in BedrockGroundedResponse.__annotations__ + assert "citations" in BedrockGroundedResponse.__annotations__ + + +def test_bedrock_empty_citations_list(bedrock_mocks): + """ + Test that responses with empty citations list return plain string. + """ + _, mock_client = bedrock_mocks + + # Response with empty citations + mock_response_empty_citations = { + 'output': { + 'message': { + 'role': 'assistant', + 'content': [ + { + 'text': 'Response with empty citations.', + 'citationsContent': { + 'citations': [] + } + } + ] + } + }, + 'usage': { + 'inputTokens': 10, + 'outputTokens': 5, + 'totalTokens': 15 + } + } + mock_client.converse.return_value = mock_response_empty_citations + + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding"], + ) + + result = llm.call("Test query") + + # Empty citations should return plain string + assert isinstance(result, str), f"Expected str for empty citations, got {type(result)}" + assert result == "Response with empty citations." + + +def test_bedrock_multiple_system_tools(): + """ + Test that multiple system tools can be configured. + """ + llm = LLM( + model="bedrock/amazon.nova-pro-v1:0", + system_tools=["nova_grounding", "another_tool"], + ) + + from crewai.llms.providers.bedrock.completion import BedrockCompletion + assert isinstance(llm, BedrockCompletion) + assert llm.system_tools == ["nova_grounding", "another_tool"] + assert len(llm.system_tools) == 2