diff --git a/README.md b/README.md index 938206d4..07828b96 100644 --- a/README.md +++ b/README.md @@ -1346,11 +1346,10 @@ The `graph` tool uses the same model provider environment variables as `use_agen | Environment Variable | Description | Default | |----------------------|-------------|---------| | MONGODB_ATLAS_CLUSTER_URI | MongoDB Atlas connection string | None | -| MONGODB_DEFAULT_DATABASE | Default database name for MongoDB operations | memories | -| MONGODB_DEFAULT_COLLECTION | Default collection name for MongoDB operations | user_memories | -| MONGODB_DEFAULT_NAMESPACE | Default namespace for memory isolation | default | -| MONGODB_DEFAULT_MAX_RESULTS | Default maximum results for list operations | 50 | -| MONGODB_DEFAULT_MIN_SCORE | Default minimum relevance score for filtering results | 0.4 | +| MONGODB_DATABASE_NAME | Database name for MongoDB operations | strands_memory | +| MONGODB_COLLECTION_NAME | Collection name for MongoDB operations | memories | +| MONGODB_NAMESPACE | Namespace for memory isolation | default | +| MONGODB_EMBEDDING_MODEL | Amazon Bedrock model for embeddings | amazon.titan-embed-text-v2:0 | **Note**: This tool requires AWS account credentials to generate embeddings using Amazon Bedrock Titan models. diff --git a/docs/mongodb_memory_tool.md b/docs/mongodb_memory_tool.md index e8d8fea2..a008cbaa 100644 --- a/docs/mongodb_memory_tool.md +++ b/docs/mongodb_memory_tool.md @@ -112,7 +112,7 @@ result = mongodb_memory( ### Environment Variables -You can also use environment variables for configuration: +You can use environment variables for configuration: ```bash export MONGODB_ATLAS_CLUSTER_URI="mongodb+srv://user:password@cluster.mongodb.net/" @@ -123,6 +123,10 @@ export MONGODB_EMBEDDING_MODEL="amazon.titan-embed-text-v2:0" export AWS_REGION="us-west-2" ``` +**Note:** Environment variables take precedence over tool parameters in the standalone function. +To let the agent control the connection target, use the class-based approach or leave the +environment variable unset. + Then use the tool with minimal parameters (environment variables will be used): ```python diff --git a/src/strands_tools/code_interpreter/agent_core_code_interpreter.py b/src/strands_tools/code_interpreter/agent_core_code_interpreter.py index 4e5ebf42..0457e639 100644 --- a/src/strands_tools/code_interpreter/agent_core_code_interpreter.py +++ b/src/strands_tools/code_interpreter/agent_core_code_interpreter.py @@ -479,7 +479,12 @@ def write_files(self, action: WriteFilesAction) -> Dict[str, Any]: logger.debug(f"Writing {len(action.content)} files to session '{session_name}'") - content_dicts = [{"path": fc.path, "text": fc.text} for fc in action.content] + content_dicts = [] + for fc in action.content: + if fc.blob is not None: + content_dicts.append({"path": fc.path, "blob": fc.blob}) + else: + content_dicts.append({"path": fc.path, "text": fc.text}) params = {"content": content_dicts} response = self._sessions[session_name].client.invoke("writeFiles", params) diff --git a/src/strands_tools/code_interpreter/models.py b/src/strands_tools/code_interpreter/models.py index b02f9b85..b826546b 100644 --- a/src/strands_tools/code_interpreter/models.py +++ b/src/strands_tools/code_interpreter/models.py @@ -8,7 +8,7 @@ from enum import Enum from typing import List, Literal, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator class LanguageType(str, Enum): @@ -24,7 +24,22 @@ class FileContent(BaseModel): updating files during code execution sessions.""" path: str = Field(description="The file path where content should be written") - text: str = Field(description="Text content for the file") + text: Optional[str] = Field( + default=None, + description="Text content for the file", + ) + blob: Optional[bytes] = Field( + default=None, + description="Base64-encoded binary content for the file", + ) + + @model_validator(mode="after") + def validate_content(self) -> "FileContent": + if self.text is None and self.blob is None: + raise ValueError("Either text or blob must be provided") + if self.text is not None and self.blob is not None: + raise ValueError("Only one of text or blob may be provided") + return self # Action-specific Pydantic models using discriminated unions diff --git a/src/strands_tools/elasticsearch_memory.py b/src/strands_tools/elasticsearch_memory.py index 5ad68f95..534f733b 100644 --- a/src/strands_tools/elasticsearch_memory.py +++ b/src/strands_tools/elasticsearch_memory.py @@ -36,7 +36,7 @@ from strands import Agent from strands_tools.elasticsearch_memory import elasticsearch_memory -# Create agent with direct tool usage +# Create agent with elasticsearch_memory tool (credentials via env vars) agent = Agent(tools=[elasticsearch_memory]) # Store a memory with semantic embeddings @@ -44,8 +44,6 @@ action="record", content="User prefers vegetarian pizza with extra cheese", metadata={"category": "food_preferences", "type": "dietary"}, - cloud_id="your-elasticsearch-cloud-id", - api_key="your-api-key", index_name="memories", namespace="user_123" ) @@ -55,8 +53,6 @@ action="retrieve", query="food preferences and dietary restrictions", max_results=5, - cloud_id="your-elasticsearch-cloud-id", - api_key="your-api-key", index_name="memories", namespace="user_123" ) @@ -65,8 +61,6 @@ elasticsearch_memory( action="list", max_results=10, - cloud_id="your-elasticsearch-cloud-id", - api_key="your-api-key", index_name="memories", namespace="user_123" ) @@ -75,8 +69,6 @@ elasticsearch_memory( action="get", memory_id="mem_1234567890_abcd1234", - cloud_id="your-elasticsearch-cloud-id", - api_key="your-api-key", index_name="memories", namespace="user_123" ) @@ -85,8 +77,6 @@ elasticsearch_memory( action="delete", memory_id="mem_1234567890_abcd1234", - cloud_id="your-elasticsearch-cloud-id", - api_key="your-api-key", index_name="memories", namespace="user_123" ) @@ -599,9 +589,6 @@ def elasticsearch_memory( max_results: Optional[int] = None, next_token: Optional[str] = None, metadata: Optional[Dict] = None, - cloud_id: Optional[str] = None, - api_key: Optional[str] = None, - es_url: Optional[str] = None, index_name: Optional[str] = None, namespace: Optional[str] = None, embedding_model: Optional[str] = None, @@ -639,6 +626,10 @@ def elasticsearch_memory( - delete: Remove a specific memory Use this to delete memories that are no longer needed. + Connection credentials are read from environment variables: + - ELASTICSEARCH_CLOUD_ID or ELASTICSEARCH_URL for connection + - ELASTICSEARCH_API_KEY for authentication + Args: action: The memory operation to perform (one of: "record", "retrieve", "list", "get", "delete") content: For record action: Text content to store as a memory @@ -647,9 +638,6 @@ def elasticsearch_memory( max_results: Maximum number of results to return (optional, default: 10) next_token: Pagination token for list action (optional) metadata: Additional metadata to store with the memory (optional) - cloud_id: Elasticsearch Cloud ID for connection (optional if es_url provided) - api_key: Elasticsearch API key for authentication - es_url: Elasticsearch URL for serverless connection (optional if cloud_id provided) index_name: Name of the Elasticsearch index (defaults to 'strands_memory') namespace: Namespace for memory operations (defaults to 'default') embedding_model: Amazon Bedrock model for embeddings (defaults to Titan) @@ -658,12 +646,11 @@ def elasticsearch_memory( Returns: Dict: Response containing the requested memory information or operation status """ - try: - # Get values from environment variables if not provided - cloud_id = cloud_id or os.getenv("ELASTICSEARCH_CLOUD_ID") - es_url = es_url or os.getenv("ELASTICSEARCH_URL") - api_key = api_key or os.getenv("ELASTICSEARCH_API_KEY") + cloud_id = os.getenv("ELASTICSEARCH_CLOUD_ID") + es_url = os.getenv("ELASTICSEARCH_URL") + api_key = os.getenv("ELASTICSEARCH_API_KEY") + try: # Validate required parameters if not api_key: return {"status": "error", "content": [{"text": "api_key is required"}]} diff --git a/src/strands_tools/environment.py b/src/strands_tools/environment.py index ab150e12..5b66a639 100644 --- a/src/strands_tools/environment.py +++ b/src/strands_tools/environment.py @@ -147,7 +147,16 @@ # Protected variables that can't be modified -PROTECTED_VARS = {"PATH", "PYTHONPATH", "STRANDS_HOME", "SHELL", "USER", "HOME", "BYPASS_TOOL_CONSENT"} +PROTECTED_VARS = { + "PATH", + "PYTHONPATH", + "STRANDS_HOME", + "SHELL", + "USER", + "HOME", + "BYPASS_TOOL_CONSENT", + "STRANDS_NON_INTERACTIVE", +} def mask_sensitive_value(name: str, value: str) -> str: diff --git a/src/strands_tools/file_read.py b/src/strands_tools/file_read.py index d57143db..094b5438 100644 --- a/src/strands_tools/file_read.py +++ b/src/strands_tools/file_read.py @@ -384,7 +384,9 @@ def find_files(console: Console, pattern: str, recursive: bool = True) -> List[s elif os.path.isdir(pattern): matching_files = [] - for root, _dirs, files in os.walk(pattern): + for root, dirs, files in os.walk(pattern): + dirs[:] = [d for d in dirs if not d.startswith(".")] + if not recursive and root != pattern: continue diff --git a/src/strands_tools/mongodb_memory.py b/src/strands_tools/mongodb_memory.py index a6877f9b..c8f70d0e 100644 --- a/src/strands_tools/mongodb_memory.py +++ b/src/strands_tools/mongodb_memory.py @@ -1006,9 +1006,12 @@ def mongodb_memory( max_results: Maximum number of results to return (optional, default: 10) next_token: Pagination token for list action (optional) metadata: Additional metadata to store with the memory (optional) - cluster_uri: MongoDB Atlas cluster URI (optional if set via environment) - database_name: Name of the MongoDB database (optional, defaults to 'strands_memory') - collection_name: Name of the MongoDB collection (optional, defaults to 'memories') + cluster_uri: MongoDB Atlas cluster URI. If the MONGODB_ATLAS_CLUSTER_URI environment variable + is set, it takes precedence and this parameter is ignored. + database_name: Name of the MongoDB database. If the MONGODB_DATABASE_NAME environment variable + is set, it takes precedence. Defaults to 'strands_memory'. + collection_name: Name of the MongoDB collection. If the MONGODB_COLLECTION_NAME environment + variable is set, it takes precedence. Defaults to 'memories'. namespace: Namespace for memory operations (defaults to 'default') embedding_model: Amazon Bedrock model for embeddings (defaults to Titan) region: AWS region for Bedrock service (defaults to 'us-west-2') @@ -1018,12 +1021,13 @@ def mongodb_memory( Dict: Response containing the requested memory information or operation status """ try: - # Get values from environment variables if not provided - cluster_uri = cluster_uri or os.getenv("MONGODB_ATLAS_CLUSTER_URI") - database_name = database_name or os.getenv("MONGODB_DATABASE_NAME", DEFAULT_DATABASE_NAME) - collection_name = collection_name or os.getenv("MONGODB_COLLECTION_NAME", DEFAULT_COLLECTION_NAME) - embedding_model = embedding_model or os.getenv("MONGODB_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL) - region = region or os.getenv("AWS_REGION", DEFAULT_AWS_REGION) + # Environment variables take precedence over agent-provided parameters to prevent + # the agent from redirecting connections to untrusted servers. + cluster_uri = os.getenv("MONGODB_ATLAS_CLUSTER_URI") or cluster_uri + database_name = os.getenv("MONGODB_DATABASE_NAME", database_name or DEFAULT_DATABASE_NAME) + collection_name = os.getenv("MONGODB_COLLECTION_NAME", collection_name or DEFAULT_COLLECTION_NAME) + embedding_model = os.getenv("MONGODB_EMBEDDING_MODEL", embedding_model or DEFAULT_EMBEDDING_MODEL) + region = os.getenv("AWS_REGION", region or DEFAULT_AWS_REGION) vector_index_name = vector_index_name or DEFAULT_VECTOR_INDEX_NAME if namespace is None: namespace = os.getenv("MONGODB_NAMESPACE", DEFAULT_NAMESPACE) diff --git a/src/strands_tools/retrieve.py b/src/strands_tools/retrieve.py index f5882d79..842be4b9 100644 --- a/src/strands_tools/retrieve.py +++ b/src/strands_tools/retrieve.py @@ -205,6 +205,20 @@ def filter_results_by_score(results: List[Dict[str, Any]], min_score: float) -> return [result for result in results if result.get("score", 0.0) >= min_score] +# Mapping of RetrievalResultLocation types to their document identifier fields. +# See: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrievalResultLocation.html +_LOCATION_FIELD_MAP = { + "customDocumentLocation": "id", + "s3Location": "uri", + "webLocation": "url", + "confluenceLocation": "url", + "salesforceLocation": "url", + "sharePointLocation": "url", + "kendraDocumentLocation": "uri", + "sqlLocation": "query", +} + + def format_results_for_display(results: List[Dict[str, Any]], enable_metadata: bool = False) -> str: """ Format retrieval results for readable display. @@ -226,14 +240,13 @@ def format_results_for_display(results: List[Dict[str, Any]], enable_metadata: b formatted = [] for result in results: - # Extract document location - handle both s3Location and customDocumentLocation + # Extract document location - handle all RetrievalResultLocation types location = result.get("location", {}) doc_id = "Unknown" - if "customDocumentLocation" in location: - doc_id = location["customDocumentLocation"].get("id", "Unknown") - elif "s3Location" in location: - # Extract meaningful part from S3 URI - doc_id = location["s3Location"].get("uri", "") + for loc_key, field in _LOCATION_FIELD_MAP.items(): + if loc_key in location: + doc_id = location[loc_key].get(field, "Unknown") + break score = result.get("score", 0.0) formatted.append(f"\nScore: {score:.4f}") formatted.append(f"Document ID: {doc_id}") diff --git a/src/strands_tools/shell.py b/src/strands_tools/shell.py index 61fe00f0..c7758770 100644 --- a/src/strands_tools/shell.py +++ b/src/strands_tools/shell.py @@ -416,7 +416,6 @@ def shell( ignore_errors: bool = False, timeout: int = None, work_dir: str = None, - non_interactive: bool = False, ) -> Dict[str, Any]: """Interactive shell with PTY support for real-time command execution and interaction. Features: @@ -482,16 +481,13 @@ def shell( ignore_errors: Continue execution even if some commands fail (default: False) timeout: Timeout in seconds for each command (default: controlled by SHELL_DEFAULT_TIMEOUT environment variable) work_dir: Working directory for command execution (default: current) - non_interactive: Run in non-interactive mode without user prompts (default: False) Returns: Dict containing status and response content """ console = console_util.create() - is_strands_non_interactive = os.environ.get("STRANDS_NON_INTERACTIVE", "").lower() == "true" - # Here we keep both doors open, but we only prompt env STRANDS_NON_INTERACTIVE in our doc. - non_interactive_mode = is_strands_non_interactive or non_interactive + non_interactive_mode = os.environ.get("STRANDS_NON_INTERACTIVE", "").lower() == "true" # Validate command parameter if command is None: diff --git a/src/strands_tools/workflow.py b/src/strands_tools/workflow.py index c9dd0555..75b829cb 100644 --- a/src/strands_tools/workflow.py +++ b/src/strands_tools/workflow.py @@ -373,6 +373,7 @@ def _create_task_agent(self, task: Dict) -> Agent: # Create the task agent task_agent = Agent( + name=task.get("task_id"), model=selected_model, system_prompt=system_prompt, tools=filtered_tools, diff --git a/tests/code_interpreter/test_agent_core_code_interpreter.py b/tests/code_interpreter/test_agent_core_code_interpreter.py index 89a7bb4e..e69550a6 100644 --- a/tests/code_interpreter/test_agent_core_code_interpreter.py +++ b/tests/code_interpreter/test_agent_core_code_interpreter.py @@ -756,6 +756,37 @@ def test_write_files_success(interpreter, mock_client): ) +def test_write_files_action_with_blob(interpreter, mock_client): + """Test successful file writing with base64 blob content.""" + session_info = SessionInfo(session_id="test-session-id-123", description="Test session", client=mock_client) + interpreter._sessions["test-session"] = session_info + + action = WriteFilesAction( + type="writeFiles", + session_name="test-session", + content=[ + FileContent(path="image.png", blob=b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"), + FileContent(path="data.txt", text="Some data"), + ], + ) + + result = interpreter.write_files(action) + + assert result["status"] == "success" + mock_client.invoke.assert_called_once_with( + "writeFiles", + { + "content": [ + { + "path": "image.png", + "blob": b"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ", + }, + {"path": "data.txt", "text": "Some data"}, + ] + }, + ) + + def test_create_tool_result_with_stream(interpreter): """Test _create_tool_result with stream response.""" response = {"stream": [{"result": {"content": "Test output"}}], "isError": False} diff --git a/tests/test_elasticsearch_memory.py b/tests/test_elasticsearch_memory.py index b95b5bc8..32482219 100644 --- a/tests/test_elasticsearch_memory.py +++ b/tests/test_elasticsearch_memory.py @@ -2,6 +2,7 @@ Tests for the elasticsearch_memory tool. """ +import inspect import json import os from unittest import mock @@ -12,6 +13,29 @@ from src.strands_tools.elasticsearch_memory import elasticsearch_memory +ES_ENV_VARS = { + "ELASTICSEARCH_CLOUD_ID": "test-cloud-id", + "ELASTICSEARCH_API_KEY": "test-api-key", + "ELASTICSEARCH_INDEX_NAME": "test_index", + "ELASTICSEARCH_NAMESPACE": "test_namespace", + "AWS_REGION": "us-east-1", +} + +ES_URL_ENV_VARS = { + "ELASTICSEARCH_URL": "https://test-cluster.es.region.aws.elastic.cloud:443", + "ELASTICSEARCH_API_KEY": "test-api-key", + "ELASTICSEARCH_INDEX_NAME": "test_index", + "ELASTICSEARCH_NAMESPACE": "test_namespace", + "AWS_REGION": "us-east-1", +} + + +@pytest.fixture(autouse=True) +def set_es_env_vars(): + """Auto-set ES environment variables for all tests.""" + with mock.patch.dict(os.environ, ES_ENV_VARS): + yield + @pytest.fixture def mock_elasticsearch_client(): @@ -74,29 +98,40 @@ def agent(mock_elasticsearch_client, mock_bedrock_client): @pytest.fixture def config(): - """Configuration parameters for testing.""" + """Configuration parameters for testing (non-credential params only).""" return { - "cloud_id": "test-cloud-id", - "api_key": "test-api-key", "index_name": "test_index", "namespace": "test_namespace", - "region": "us-east-1", } +# --- Security Tests --- + + +def test_credential_params_not_in_tool_signature(): + """Verify that credential/connection parameters cannot be passed to the tool.""" + sig = inspect.signature(elasticsearch_memory) + param_names = set(sig.parameters.keys()) + assert "es_url" not in param_names + assert "cloud_id" not in param_names + assert "api_key" not in param_names + + def test_missing_required_params(mock_elasticsearch_client, mock_bedrock_client): - """Test tool with missing required parameters.""" + """Test tool with missing required environment variables.""" agent = Agent(tools=[elasticsearch_memory]) - # Test missing both cloud_id and es_url - result = agent.tool.elasticsearch_memory(action="record", content="test", api_key="test-api-key") - assert result["status"] == "error" - assert "Either cloud_id or es_url is required" in result["content"][0]["text"] + # Test missing both cloud_id and es_url (no env vars set) + with mock.patch.dict(os.environ, {"ELASTICSEARCH_API_KEY": "test-api-key"}, clear=True): + result = agent.tool.elasticsearch_memory(action="record", content="test") + assert result["status"] == "error" + assert "Either cloud_id or es_url is required" in result["content"][0]["text"] # Test missing api_key - result = agent.tool.elasticsearch_memory(action="record", content="test", cloud_id="test-cloud-id") - assert result["status"] == "error" - assert "api_key is required" in result["content"][0]["text"] + with mock.patch.dict(os.environ, {"ELASTICSEARCH_CLOUD_ID": "test-cloud-id"}, clear=True): + result = agent.tool.elasticsearch_memory(action="record", content="test") + assert result["status"] == "error" + assert "api_key is required" in result["content"][0]["text"] def test_connection_failure(mock_elasticsearch_client, mock_bedrock_client): @@ -106,9 +141,8 @@ def test_connection_failure(mock_elasticsearch_client, mock_bedrock_client): # Configure ping to return False (connection failure) mock_elasticsearch_client["client"].ping.return_value = False - result = agent.tool.elasticsearch_memory( - action="record", content="test", cloud_id="test-cloud-id", api_key="test-api-key" - ) + with mock.patch.dict(os.environ, ES_ENV_VARS): + result = agent.tool.elasticsearch_memory(action="record", content="test") assert result["status"] == "error" assert "Unable to connect to Elasticsearch cluster" in result["content"][0]["text"] @@ -121,7 +155,8 @@ def test_index_creation(mock_elasticsearch_client, mock_bedrock_client, config): # Configure mock responses mock_elasticsearch_client["client"].index.return_value = {"result": "created", "_id": "test_memory_id"} - agent.tool.elasticsearch_memory(action="record", content="Test content", **config) + with mock.patch.dict(os.environ, ES_ENV_VARS): + agent.tool.elasticsearch_memory(action="record", content="Test content", **config) # Verify index creation was called mock_elasticsearch_client["client"].indices.create.assert_called_once() @@ -158,9 +193,10 @@ def test_record_memory(mock_elasticsearch_client, mock_bedrock_client, config): mock_elasticsearch_client["client"].index.return_value = {"result": "created", "_id": "test_memory_id"} # Call the tool - result = agent.tool.elasticsearch_memory( - action="record", content="Test memory content", metadata={"category": "test"}, **config - ) + with mock.patch.dict(os.environ, ES_ENV_VARS): + result = agent.tool.elasticsearch_memory( + action="record", content="Test memory content", metadata={"category": "test"}, **config + ) # Verify success response assert result["status"] == "success" @@ -504,8 +540,6 @@ def test_agent_tool_usage(mock_elasticsearch_client, mock_bedrock_client): result = agent.tool.elasticsearch_memory( action="record", content="Test memory content", - cloud_id="test-cloud-id", - api_key="test-api-key", index_name="test_index", namespace="test_namespace", ) @@ -522,21 +556,20 @@ def test_agent_tool_usage(mock_elasticsearch_client, mock_bedrock_client): def test_es_url_connection(mock_elasticsearch_client, mock_bedrock_client): - """Test using es_url instead of cloud_id for connection.""" + """Test using ELASTICSEARCH_URL instead of ELASTICSEARCH_CLOUD_ID for connection.""" agent = Agent(tools=[elasticsearch_memory]) # Configure mock responses mock_elasticsearch_client["client"].index.return_value = {"result": "created", "_id": "test_memory_id"} - # Call tool with es_url instead of cloud_id - result = agent.tool.elasticsearch_memory( - action="record", - content="Test memory content", - es_url="https://test-cluster.es.region.aws.elastic.cloud:443", - api_key="test-api-key", - index_name="test_index", - namespace="test_namespace", - ) + # Override env vars to use URL instead of cloud_id + with mock.patch.dict(os.environ, ES_URL_ENV_VARS, clear=False): + result = agent.tool.elasticsearch_memory( + action="record", + content="Test memory content", + index_name="test_index", + namespace="test_namespace", + ) # Verify success response assert result["status"] == "success" @@ -558,7 +591,8 @@ def test_custom_embedding_model(mock_elasticsearch_client, mock_bedrock_client, # Call tool with custom embedding model result = agent.tool.elasticsearch_memory( - action="record", content="Test memory content", embedding_model="amazon.titan-embed-text-v1:0", **config + action="record", content="Test memory content", embedding_model="amazon.titan-embed-text-v1:0", + index_name="test_index", namespace="test_namespace", ) # Verify success response @@ -583,7 +617,7 @@ def test_multiple_namespaces(mock_elasticsearch_client, mock_bedrock_client, con action="record", content="Alice likes Italian food", namespace="user_alice", - **{k: v for k, v in config.items() if k != "namespace"}, + index_name="test_index", ) # Store memory in system namespace @@ -591,7 +625,7 @@ def test_multiple_namespaces(mock_elasticsearch_client, mock_bedrock_client, con action="record", content="System maintenance scheduled", namespace="system_global", - **{k: v for k, v in config.items() if k != "namespace"}, + index_name="test_index", ) # Verify both operations succeeded @@ -626,13 +660,10 @@ def test_configuration_dictionary_pattern(mock_elasticsearch_client, mock_bedroc } } - # Create configuration dictionary + # Create configuration dictionary (non-credential params only) config = { - "cloud_id": "test-cloud-id", - "api_key": "test-api-key", "index_name": "memories", "namespace": "user_123", - "region": "us-east-1", } # Store memory using config dictionary @@ -782,14 +813,13 @@ def test_security_scenarios(mock_elasticsearch_client, mock_bedrock_client): } # Test namespace validation - memory exists but not in requested namespace - result = agent.tool.elasticsearch_memory( - action="get", - memory_id="mem_123", - cloud_id="test-cloud-id", - api_key="test-api-key", - index_name="test_index", - namespace="correct_namespace", - ) + with mock.patch.dict(os.environ, ES_ENV_VARS): + result = agent.tool.elasticsearch_memory( + action="get", + memory_id="mem_123", + index_name="test_index", + namespace="correct_namespace", + ) assert result["status"] == "error" assert "not found in namespace correct_namespace" in result["content"][0]["text"] @@ -819,8 +849,8 @@ def test_injection_prevention(mock_elasticsearch_client, mock_bedrock_client, co """Test that injection attempts via namespace are blocked.""" agent = Agent(tools=[elasticsearch_memory]) - # Remove namespace from config to avoid conflict - test_config = {k: v for k, v in config.items() if k != "namespace"} + # Config without namespace + test_config = {"index_name": "test_index"} # Test dict-based injection (analogous to MongoDB {"$ne": ""} attack) malicious_namespace = {"$ne": ""} @@ -848,8 +878,8 @@ def test_namespace_validation_strict_rules(mock_elasticsearch_client, mock_bedro """Test strict namespace validation rules.""" agent = Agent(tools=[elasticsearch_memory]) - # Remove namespace from config to avoid conflict - test_config = {k: v for k, v in config.items() if k != "namespace"} + # Config without namespace + test_config = {"index_name": "test_index"} # Test invalid characters (should be rejected) invalid_namespaces = [ @@ -883,8 +913,8 @@ def test_valid_namespaces_accepted(mock_elasticsearch_client, mock_bedrock_clien } } - # Remove namespace from config - test_config = {k: v for k, v in config.items() if k != "namespace"} + # Config without namespace + test_config = {"index_name": "test_index"} valid_namespaces = [ "default", diff --git a/tests/test_file_read.py b/tests/test_file_read.py index c0db8e35..b05a1262 100644 --- a/tests/test_file_read.py +++ b/tests/test_file_read.py @@ -280,6 +280,45 @@ def test_find_files_function(): assert len(files) == 2 # Should not find test3.txt in subdir +def test_find_files_skips_hidden_directories(): + """Test that find_files skips hidden directories during traversal.""" + mock_console = unittest.mock.Mock() + + with tempfile.TemporaryDirectory() as temp_dir: + # Create normal files + normal_file = os.path.join(temp_dir, "main.py") + with open(normal_file, "w") as f: + f.write("normal") + + # Create hidden directory with files (simulates .venv, .git, etc.) + hidden_dir = os.path.join(temp_dir, ".hidden_dir") + os.makedirs(hidden_dir) + hidden_file = os.path.join(hidden_dir, "should_not_appear.py") + with open(hidden_file, "w") as f: + f.write("hidden") + + # Create nested hidden directory + nested_hidden = os.path.join(temp_dir, "subdir", ".cache") + os.makedirs(nested_hidden) + nested_hidden_file = os.path.join(nested_hidden, "cached.pyc") + with open(nested_hidden_file, "w") as f: + f.write("cached") + + # Create normal subdirectory + normal_sub = os.path.join(temp_dir, "subdir") + normal_sub_file = os.path.join(normal_sub, "utils.py") + with open(normal_sub_file, "w") as f: + f.write("utils") + + files = file_read.find_files(mock_console, temp_dir, recursive=True) + + filenames = [os.path.basename(f) for f in files] + assert "main.py" in filenames + assert "utils.py" in filenames + assert "should_not_appear.py" not in filenames + assert "cached.pyc" not in filenames + + def test_split_path_list_function(): """Test the split_path_list utility function.""" paths = "path1.txt,path2.md, path3.py" diff --git a/tests/test_retrieve.py b/tests/test_retrieve.py index 97add221..fa9968ca 100644 --- a/tests/test_retrieve.py +++ b/tests/test_retrieve.py @@ -134,6 +134,75 @@ def test_format_results_for_display(): assert "Content: S3 content" in s3_formatted +@pytest.mark.parametrize( + "location_key,location_data,location_type,expected_doc_id", + [ + ("webLocation", {"url": "https://example.com/docs/page.html"}, "WEB", "https://example.com/docs/page.html"), + ( + "confluenceLocation", + {"url": "https://mycompany.atlassian.net/wiki/spaces/DOC/pages/123"}, + "CONFLUENCE", + "https://mycompany.atlassian.net/wiki/spaces/DOC/pages/123", + ), + ( + "salesforceLocation", + {"url": "https://mycompany.salesforce.com/articles/KB001"}, + "SALESFORCE", + "https://mycompany.salesforce.com/articles/KB001", + ), + ( + "sharePointLocation", + {"url": "https://mycompany.sharepoint.com/sites/docs/page.aspx"}, + "SHAREPOINT", + "https://mycompany.sharepoint.com/sites/docs/page.aspx", + ), + ( + "kendraDocumentLocation", + {"uri": "https://kendra.aws/documents/doc-12345"}, + "KENDRA", + "https://kendra.aws/documents/doc-12345", + ), + ( + "sqlLocation", + {"query": "SELECT * FROM documents WHERE id = 1"}, + "SQL", + "SELECT * FROM documents WHERE id = 1", + ), + ], +) +def test_format_results_for_display_location_types(location_key, location_data, location_type, expected_doc_id): + """Test format_results_for_display with all supported RetrievalResultLocation types.""" + test_results = [ + { + "content": {"text": "Test content", "type": "TEXT"}, + "location": {location_key: location_data, "type": location_type}, + "score": 0.80, + } + ] + formatted = retrieve.format_results_for_display(test_results) + assert f"Document ID: {expected_doc_id}" in formatted + assert "Content: Test content" in formatted + + +def test_format_results_for_display_unknown_location(): + """Test format_results_for_display with an unrecognized location type.""" + test_results = [ + { + "content": {"text": "Unknown source content", "type": "TEXT"}, + "location": { + "futureLocation": {"url": "https://future.example.com"}, + "type": "FUTURE", + }, + "score": 0.70, + } + ] + + formatted = retrieve.format_results_for_display(test_results) + assert "Score: 0.7000" in formatted + assert "Document ID: Unknown" in formatted + assert "Content: Unknown source content" in formatted + + def test_format_results_with_metadata(): """Test the format_results_for_display function with metadata enabled.""" test_results = [ diff --git a/tests/test_shell.py b/tests/test_shell.py index 5c253eb1..8d3f20e7 100644 --- a/tests/test_shell.py +++ b/tests/test_shell.py @@ -823,8 +823,9 @@ def test_shell_with_json_array_string(mock_execute_commands): "status": "success", }, ] - # Call the shell function with new @tool decorator signature and non_interactive=True to skip confirmation - shell.shell('["cmd1", "cmd2"]', non_interactive=True) + # Call the shell function with STRANDS_NON_INTERACTIVE env var to skip confirmation + with patch.dict(os.environ, {"STRANDS_NON_INTERACTIVE": "true"}): + shell.shell('["cmd1", "cmd2"]') # Verify execute_commands was called with parsed array mock_execute_commands.assert_called_once() @@ -846,8 +847,9 @@ def test_shell_with_invalid_json_array_string(mock_execute_commands): } ] - # Call the shell function with new @tool decorator signature and non_interactive=True - shell.shell("[invalid json array]", non_interactive=True) + # Call the shell function with STRANDS_NON_INTERACTIVE env var to skip confirmation + with patch.dict(os.environ, {"STRANDS_NON_INTERACTIVE": "true"}): + shell.shell("[invalid json array]") # Verify execute_commands was called with the original string (fallback) mock_execute_commands.assert_called_once() @@ -870,9 +872,10 @@ def test_shell_with_complex_command_objects(mock_execute_commands): } ] - # Call the shell function with new @tool decorator signature - complex command object + # Call the shell function with STRANDS_NON_INTERACTIVE env var - complex command object command_obj = {"command": "git clone", "timeout": 120} - shell.shell(command_obj, non_interactive=True) + with patch.dict(os.environ, {"STRANDS_NON_INTERACTIVE": "true"}): + shell.shell(command_obj) # Verify execute_commands was called with the complex object mock_execute_commands.assert_called_once() @@ -936,8 +939,9 @@ def test_shell_console_output(mock_console_util): # Reset mock for non-interactive test mock_console.reset_mock() - # Call with non_interactive_mode=True - shell.shell("echo test", non_interactive=True) + # Call with STRANDS_NON_INTERACTIVE env var + with patch.dict(os.environ, {"STRANDS_NON_INTERACTIVE": "true"}): + shell.shell("echo test") # Verify console.print was not called for UI elements in non-interactive mode assert mock_console.print.call_count == 0 @@ -967,32 +971,11 @@ def test_shell_exception_ui_output(mock_execute_commands, mock_get_user_input, m # Reset mock for non-interactive test mock_console.reset_mock() - # Call with non_interactive_mode=True - shell.shell("echo test", non_interactive=True) + # Call with STRANDS_NON_INTERACTIVE env var + with patch.dict(os.environ, {"STRANDS_NON_INTERACTIVE": "true"}): + shell.shell("echo test") # Verify console.print was not called in non-interactive mode assert mock_console.print.call_count == 0 -@patch("strands_tools.shell.execute_commands") -@patch("strands_tools.shell.get_user_input") -def test_shell_non_interactive_parameter(mock_get_user_input, mock_execute_commands): - """Test shell tool with non_interactive parameter.""" - # Mock execute_commands to return a successful result - mock_execute_commands.return_value = [ - { - "command": "echo test", - "exit_code": 0, - "output": "test\n", - "error": "", - "status": "success", - } - ] - - # Call the shell function with non_interactive=True - result = shell.shell("echo test", non_interactive=True) - - assert result["status"] == "success" - - # Verify that get_user_input was not called because non_interactive=True - mock_get_user_input.assert_not_called()