Add attachment search tool helper

src/xai_sdk/tools.py

@@ -166,6 +166,33 @@ def code_execution() -> chat_pb2.Tool:
     return chat_pb2.Tool(code_execution=chat_pb2.CodeExecution())
 
 
+def attachment_search(limit: Optional[int] = None) -> chat_pb2.Tool:
+    """Creates a server-side tool for searching files attached to the request.
+
+    This tool enables the model to search over file attachments provided in the current
+    request. It can be configured with an optional maximum number of results.
+
+    Args:
+        limit: The maximum number of attachment search results to return. If omitted,
+            the server uses its default limit.
+
+    Returns:
+        A `chat_pb2.Tool` object configured for attachment search.
+
+    Example:
+        ```
+        from xai_sdk.tools import attachment_search
+
+        tool = attachment_search(limit=5)
+        ```
+    """
+    attachment_search_kwargs: dict = {}
+    if limit is not None:
+        attachment_search_kwargs["limit"] = limit
+
+    return chat_pb2.Tool(attachment_search=chat_pb2.AttachmentSearch(**attachment_search_kwargs))
+
+
 def collections_search(
     collection_ids: list[str],
     limit: Optional[int] = None,

tests/chat_test.py

@@ -645,6 +645,27 @@ def test_web_search_user_location():
     assert tool.web_search.user_location.timezone == "America/Los_Angeles"
 
 
+def test_attachment_search_tool():
+    """Test that attachment_search util function creates the expected tool."""
+    from xai_sdk.tools import attachment_search
+
+    tool = attachment_search()
+
+    assert isinstance(tool, chat_pb2.Tool)
+    assert tool.HasField("attachment_search")
+    assert not tool.attachment_search.HasField("limit")
+
+
+def test_attachment_search_tool_with_limit():
+    """Test that attachment_search util function sets the optional limit."""
+    from xai_sdk.tools import attachment_search
+
+    tool = attachment_search(limit=5)
+
+    assert tool.HasField("attachment_search")
+    assert tool.attachment_search.limit == 5
+
+
 def test_developer_message():
     """Test that developer() creates a message with ROLE_DEVELOPER role."""
     # Simple string content