feat: Search with file type - #508
Conversation
📝 WalkthroughWalkthroughAdds MIME type validation and a friendly-name → Google Drive MIME mapping in gdrive helpers, and uses it to add an optional file_type filter to search_drive_files() and list_drive_items(); extensive tests cover aliases, raw MIME inputs, query composition, and error cases. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant DriveTools as "search_drive_files / list_drive_items"
participant MimeResolver as "resolve_file_type_mime"
participant GoogleDrive as "Google Drive API"
Client->>DriveTools: Call with optional file_type
alt file_type provided
DriveTools->>MimeResolver: resolve_file_type_mime(file_type)
alt Friendly name / alias
MimeResolver->>MimeResolver: lookup FILE_TYPE_MIME_MAP (case-insensitive)
MimeResolver-->>DriveTools: return mapped MIME type
else Raw MIME (contains '/')
MimeResolver->>MimeResolver: validate against MIME_TYPE_PATTERN
MimeResolver-->>DriveTools: return raw MIME if valid
end
alt unknown/invalid
MimeResolver-->>DriveTools: raise ValueError
end
end
DriveTools->>DriveTools: append "mimeType = '...'" to query when resolved
DriveTools->>GoogleDrive: execute query (with mimeType filter if set)
GoogleDrive-->>DriveTools: return results
DriveTools-->>Client: return filtered file list or error
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@gdrive/drive_helpers.py`:
- Around line 259-267: Trim whitespace from file_type, and instead of returning
any string containing "/", validate it against a safe MIME token pattern before
returning; e.g., in the branch that currently checks "if '/' in file_type"
(using the existing variable file_type and FILE_TYPE_MIME_MAP), first do
file_type = file_type.strip(), then verify it matches a strict MIME-type regex
(token/token with allowed chars and no quotes, semicolons or control chars) and
only return it if it passes; otherwise raise a ValueError explaining the
expected MIME format or to use a friendly name.
In `@gdrive/drive_tools.py`:
- Around line 78-80: Update the docstring for the file_type parameter in
gdrive/drive_tools.py to reflect the full set of supported friendly type aliases
used in gdrive/drive_helpers.py (add 'script', 'site', 'jam', 'jamboard'
alongside existing examples like 'folder', 'document', 'doc', 'spreadsheet',
'sheet', 'presentation', 'slides', 'form', 'drawing', 'pdf', 'shortcut'), and
note that raw MIME type strings (e.g. 'application/pdf') are also accepted;
ensure the same correction is applied to the other occurrence referenced (lines
~441-445) so both docstrings accurately match DriveHelpers' supported aliases.
- Around line 106-109: When adding the MIME filter in the search logic (where
file_type is checked and resolve_file_type_mime is called), wrap the existing
structured query in parentheses before appending the "and mimeType = '...'"
clause to preserve OR semantics; e.g., set final_query = f"({final_query})" (if
not already parenthesized) and then append f" and mimeType = '{mime}'" and keep
the logger.info call (logger.info(f"[search_drive_files] Added mimeType filter:
'{mime}'")) unchanged. Ensure this change is applied in the same block that
references file_type, mime, resolve_file_type_mime, and final_query.
In `@tests/gdrive/test_drive_tools.py`:
- Around line 221-237: Add a new assertion in
test_search_file_type_structured_query_combined to cover an existing
structured-query OR case: call search_drive_files with query "name contains 'a'
or name contains 'b'" (via _unwrap) and file_type="spreadsheet", then extract
the built q from mock_service.files.return_value.list.call_args.kwargs and
assert that the original OR is preserved and that the appended mimeType filter
is grouped (e.g., parentheses around the OR clause or around the appended
mimeType as appropriate) to avoid precedence regressions; update the test name
or add a separate test function if needed to reflect this OR-case coverage.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/gdrive/test_drive_tools.py (1)
669-911: Consider parametrizing repetitivefile_typealias tests.Most cases differ only by input alias and expected MIME; a
pytest.mark.parametrizetable would reduce duplication and simplify future alias additions.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/gdrive/test_drive_tools.py` around lines 669 - 911, These repetitive tests for file_type aliases should be collapsed into parametrized tests: replace separate tests like test_search_file_type_document_alias, test_search_file_type_sheet_alias, test_search_file_type_plural_alias, test_search_file_type_folder_adds_mime_filter, test_search_file_type_raw_mime (and their list_drive_items equivalents) with pytest.mark.parametrize tables that pass (alias, expected_mime, query, optional_expected_name) into a single test for search_drive_files and a single test for list_drive_items; use the existing mocks (mock_service and patched resolve_folder_id) and assert that call_kwargs = mock_service.files.return_value.list.call_args.kwargs contains the expected "mimeType = '<expected_mime>'" (or does not contain "mimeType" for None) and preserve checks for returned text when expected_name is provided.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@gdrive/drive_helpers.py`:
- Around line 233-238: Update the stale comments that claim raw MIME values are
accepted/returned “as-is” to reflect the current behavior: raw MIME inputs are
validated against MIME_TYPE_PATTERN and normalized to lowercase before
use/return; reference the MIME_TYPE_PATTERN constant and places that handle raw
inputs (e.g., the code paths that consult FILE_TYPE_MIME_MAP and the logic
around normalization/validation around the later raw-MIME handling) and change
wording to explicitly state “validated against MIME_TYPE_PATTERN and lowercased”
instead of “as-is.”
---
Nitpick comments:
In `@tests/gdrive/test_drive_tools.py`:
- Around line 669-911: These repetitive tests for file_type aliases should be
collapsed into parametrized tests: replace separate tests like
test_search_file_type_document_alias, test_search_file_type_sheet_alias,
test_search_file_type_plural_alias,
test_search_file_type_folder_adds_mime_filter, test_search_file_type_raw_mime
(and their list_drive_items equivalents) with pytest.mark.parametrize tables
that pass (alias, expected_mime, query, optional_expected_name) into a single
test for search_drive_files and a single test for list_drive_items; use the
existing mocks (mock_service and patched resolve_folder_id) and assert that
call_kwargs = mock_service.files.return_value.list.call_args.kwargs contains the
expected "mimeType = '<expected_mime>'" (or does not contain "mimeType" for
None) and preserve checks for returned text when expected_name is provided.
Closes: #507
Feature described here. Allows searching and listing folders.
The server allows searching for specific file types, like docs or spreadsheets, by using the functions in each specific module. However, list_drive_items and search_drive_items don't allow for any filter. This is an issue when searching for folders (right now, the response includes tons of times from all the types).
This is a file type filter in list_drive_items and search_drive_items.
It works either by passing the mime type or by passing a human friendly word (like 'document').
Alternatives I've considered:
Another option was to add separate functions: list_drive_folders and search_drive_folders
These seem like an unnecessary duplication, though.
I've also considered that the file type could be a list of strings instead of a single string. However, this will add complexity for the LLM that uses it, and it may be so infrequent that it is not worth it.
There are some new tests included for list_drive_items and search_drive_items.
Checklist
Summary by CodeRabbit
New Features
Tests