-
Notifications
You must be signed in to change notification settings - Fork 115
Implement workspace-level search (DATAMAN-163) #441
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | |||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,42 @@ | |||||||||||||||||||||||||||||||||||||||||
| """Manual demo for workspace-level search (DATAMAN-163). | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| Usage: | |||||||||||||||||||||||||||||||||||||||||
| python tests/manual/demo_workspace_search.py | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| Uses staging credentials from CLAUDE.md. | |||||||||||||||||||||||||||||||||||||||||
| """ | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| import os | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| import roboflow | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| thisdir = os.path.dirname(os.path.abspath(__file__)) | |||||||||||||||||||||||||||||||||||||||||
| os.environ["ROBOFLOW_CONFIG_DIR"] = f"{thisdir}/data/.config" | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| WORKSPACE = "model-evaluation-workspace" | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| rf = roboflow.Roboflow() | |||||||||||||||||||||||||||||||||||||||||
| ws = rf.workspace(WORKSPACE) | |||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| # --- Single page search --- | |||||||||||||||||||||||||||||||||||||||||
| print("=== Single page search ===") | |||||||||||||||||||||||||||||||||||||||||
| page = ws.search("project:false", page_size=5) | |||||||||||||||||||||||||||||||||||||||||
| print(f"Total results: {page['total']}") | |||||||||||||||||||||||||||||||||||||||||
| print(f"Results in this page: {len(page['results'])}") | |||||||||||||||||||||||||||||||||||||||||
| print(f"Continuation token: {page.get('continuationToken')}") | |||||||||||||||||||||||||||||||||||||||||
|
|||||||||||||||||||||||||||||||||||||||||
| for img in page["results"]: | |||||||||||||||||||||||||||||||||||||||||
| print(f" - {img.get('filename', 'N/A')}") | |||||||||||||||||||||||||||||||||||||||||
|
|||||||||||||||||||||||||||||||||||||||||
|
|
|||||||||||||||||||||||||||||||||||||||||
| # --- Paginated search_all --- | |||||||||||||||||||||||||||||||||||||||||
| print("\n=== Paginated search_all (page_size=3, max 2 pages) ===") | |||||||||||||||||||||||||||||||||||||||||
| count = 0 | |||||||||||||||||||||||||||||||||||||||||
| for page_results in ws.search_all("*", page_size=3): | |||||||||||||||||||||||||||||||||||||||||
| count += 1 | |||||||||||||||||||||||||||||||||||||||||
| print(f"Page {count}: {len(page_results)} results") | |||||||||||||||||||||||||||||||||||||||||
| for img in page_results: | |||||||||||||||||||||||||||||||||||||||||
| print(f" - {img.get('filename', 'N/A')}") | |||||||||||||||||||||||||||||||||||||||||
|
|||||||||||||||||||||||||||||||||||||||||
| @@ -271,5 +271,11 @@ | ||
|
|
||
| def __str__(self): | ||
| """to string function""" | ||
| json_value = {"api_key": self.api_key, "workspace": self.workspace} | ||
| # Avoid exposing the full API key when this object is printed or logged. | ||
| api_key = self.api_key or "" | ||
| if len(api_key) > 4: | ||
| masked_api_key = ("*" * (len(api_key) - 4)) + api_key[-4:] | ||
| else: | ||
| masked_api_key = "*" * len(api_key) | ||
| json_value = {"api_key": masked_api_key, "workspace": self.workspace} | ||
| return json.dumps(json_value, indent=2) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's just for local run
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,138 @@ | ||
| import json | ||
| import unittest | ||
|
|
||
| import responses | ||
|
|
||
| from roboflow.adapters.rfapi import RoboflowError | ||
| from roboflow.config import API_URL | ||
|
|
||
|
|
||
| class TestWorkspaceSearch(unittest.TestCase): | ||
| API_KEY = "test_key" | ||
| WORKSPACE = "test-ws" | ||
| SEARCH_URL = f"{API_URL}/{WORKSPACE}/search/v1?api_key={API_KEY}" | ||
|
|
||
| def _make_workspace(self): | ||
| from roboflow.core.workspace import Workspace | ||
|
|
||
| info = { | ||
| "workspace": { | ||
| "name": "Test", | ||
| "url": self.WORKSPACE, | ||
| "projects": [], | ||
| "members": [], | ||
| } | ||
| } | ||
| return Workspace(info, api_key=self.API_KEY, default_workspace=self.WORKSPACE, model_format="yolov8") | ||
|
|
||
| # --- search() tests --- | ||
|
|
||
| @responses.activate | ||
| def test_search_basic(self): | ||
| body = { | ||
| "results": [{"filename": "a.jpg"}, {"filename": "b.jpg"}], | ||
| "total": 2, | ||
| "continuationToken": None, | ||
| } | ||
| responses.add(responses.POST, self.SEARCH_URL, json=body, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| result = ws.search("tag:review") | ||
|
|
||
| self.assertEqual(result["total"], 2) | ||
| self.assertEqual(len(result["results"]), 2) | ||
| self.assertIsNone(result["continuationToken"]) | ||
|
|
||
| # Verify request payload | ||
| sent = json.loads(responses.calls[0].request.body) | ||
| self.assertEqual(sent["query"], "tag:review") | ||
| self.assertEqual(sent["pageSize"], 50) | ||
| self.assertEqual(sent["fields"], ["tags", "projects", "filename"]) | ||
| self.assertNotIn("continuationToken", sent) | ||
|
|
||
| @responses.activate | ||
| def test_search_with_continuation_token(self): | ||
| body = {"results": [{"filename": "c.jpg"}], "total": 3, "continuationToken": None} | ||
| responses.add(responses.POST, self.SEARCH_URL, json=body, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| ws.search("*", continuation_token="tok_abc") | ||
|
|
||
| sent = json.loads(responses.calls[0].request.body) | ||
| self.assertEqual(sent["continuationToken"], "tok_abc") | ||
|
|
||
| @responses.activate | ||
| def test_search_custom_fields(self): | ||
| body = {"results": [], "total": 0, "continuationToken": None} | ||
| responses.add(responses.POST, self.SEARCH_URL, json=body, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| ws.search("*", fields=["filename", "embedding"]) | ||
|
|
||
| sent = json.loads(responses.calls[0].request.body) | ||
| self.assertEqual(sent["fields"], ["filename", "embedding"]) | ||
|
|
||
| @responses.activate | ||
| def test_search_api_error(self): | ||
| responses.add(responses.POST, self.SEARCH_URL, json={"error": "unauthorized"}, status=401) | ||
|
|
||
| ws = self._make_workspace() | ||
| with self.assertRaises(RoboflowError): | ||
| ws.search("tag:review") | ||
|
|
||
| # --- search_all() tests --- | ||
|
|
||
| @responses.activate | ||
| def test_search_all_single_page(self): | ||
| body = { | ||
| "results": [{"filename": "a.jpg"}, {"filename": "b.jpg"}], | ||
| "total": 2, | ||
| "continuationToken": None, | ||
| } | ||
| responses.add(responses.POST, self.SEARCH_URL, json=body, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| pages = list(ws.search_all("*")) | ||
|
|
||
| self.assertEqual(len(pages), 1) | ||
| self.assertEqual(len(pages[0]), 2) | ||
|
|
||
| @responses.activate | ||
| def test_search_all_multiple_pages(self): | ||
| page1 = { | ||
| "results": [{"filename": "a.jpg"}], | ||
| "total": 2, | ||
| "continuationToken": "tok_page2", | ||
| } | ||
| page2 = { | ||
| "results": [{"filename": "b.jpg"}], | ||
| "total": 2, | ||
| "continuationToken": None, | ||
| } | ||
| responses.add(responses.POST, self.SEARCH_URL, json=page1, status=200) | ||
| responses.add(responses.POST, self.SEARCH_URL, json=page2, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| pages = list(ws.search_all("*", page_size=1)) | ||
|
|
||
| self.assertEqual(len(pages), 2) | ||
| self.assertEqual(pages[0][0]["filename"], "a.jpg") | ||
| self.assertEqual(pages[1][0]["filename"], "b.jpg") | ||
|
|
||
| # Verify second request used the continuation token | ||
| sent2 = json.loads(responses.calls[1].request.body) | ||
| self.assertEqual(sent2["continuationToken"], "tok_page2") | ||
|
|
||
| @responses.activate | ||
| def test_search_all_empty_results(self): | ||
| body = {"results": [], "total": 0, "continuationToken": None} | ||
| responses.add(responses.POST, self.SEARCH_URL, json=body, status=200) | ||
|
|
||
| ws = self._make_workspace() | ||
| pages = list(ws.search_all("*")) | ||
|
|
||
| self.assertEqual(len(pages), 0) | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| unittest.main() |
Uh oh!
There was an error while loading. Please reload this page.