diff --git a/api_v1/utils/tencentcloud_tiia.py b/api_v1/utils/tencentcloud_tiia.py index 918e0fe..fe9f5d1 100644 --- a/api_v1/utils/tencentcloud_tiia.py +++ b/api_v1/utils/tencentcloud_tiia.py @@ -45,9 +45,25 @@ class SimpleRateLimiter: time.sleep(self._min_interval - elapsed) self._last_ts = time.monotonic() +def _strip_data_url_base64_prefix(s: str) -> str: + """ + Accept both raw base64 and data URL format like: + data:image/png;base64,AAAA... + Return the base64 payload part. + """ + s = (s or '').strip() + if not s: + return s + lower = s[:64].lower() + if lower.startswith('data:') and 'base64,' in lower: + return s.split('base64,', 1)[1].strip() + return s + + def search_image_url_in_tencent_tiia( *, - image_url: str, + image_url: str | None = None, + image_base64: str | None = None, limit: int | None = 10, offset: int | None = 0, match_threshold: int | None = None, @@ -69,7 +85,8 @@ def search_image_url_in_tencent_tiia( - TENCENTCLOUD_TIIA_ENDPOINT (default: tiia.tencentcloudapi.com) Args: - image_url: Absolute URL to search. + image_url: Absolute URL to search. (preferred when both provided) + image_base64: Base64-encoded image content. limit/offset/match_threshold: Optional SearchImage request parameters. rate_limiter: Optional in-process limiter to keep requests under QPS. @@ -77,12 +94,15 @@ def search_image_url_in_tencent_tiia( Parsed JSON response dict from TencentCloud SDK. """ image_url = (image_url or '').strip() - if not image_url: - raise ValueError('image_url 不能为空') + image_base64 = _strip_data_url_base64_prefix(image_base64 or '') - parsed = urlparse(image_url) - if not (parsed.scheme and parsed.netloc): - raise ValueError('image_url 必须为绝对 URL(包含 scheme 与 host)') + if not image_url and not image_base64: + raise ValueError('imageUrl 和 imageBase64 必须至少提供一个') + + if image_url: + parsed = urlparse(image_url) + if not (parsed.scheme and parsed.netloc): + raise ValueError('imageUrl 必须为绝对 URL(包含 scheme 与 host)') group_id = (getattr(settings, 'TENCENTCLOUD_TIIA_GROUP_ID', '') or '').strip() if not group_id: @@ -117,7 +137,11 @@ def search_image_url_in_tencent_tiia( req = models.SearchImageRequest() req.GroupId = str(group_id) - req.ImageUrl = str(image_url) + # TencentCloud: ImageUrl + ImageBase64 can be both provided, but ImageUrl wins. + if image_url: + req.ImageUrl = str(image_url) + else: + req.ImageBase64 = str(image_base64) if limit is not None: req.Limit = int(limit) diff --git a/api_v1/views/test_tiia_search_image_api.py b/api_v1/views/test_tiia_search_image_api.py index 1fa8a3b..759231b 100644 --- a/api_v1/views/test_tiia_search_image_api.py +++ b/api_v1/views/test_tiia_search_image_api.py @@ -41,6 +41,37 @@ class TiiaSearchImageAPITestCase(TestCase): ) self.assertEqual(resp.status_code, status.HTTP_200_OK) self.assertEqual(resp.json(), {'Candidates': [{'EntityId': '1', 'Score': 99}]}) + mock_search.assert_called() + + @patch('api_v1.views.tiia.search_image_url_in_tencent_tiia') + def test_search_image_success_base64(self, mock_search): + self.client.force_authenticate(self.user) + mock_search.return_value = {'Candidates': []} + + resp = self.client.post( + '/api/v1/tiia/search-image/', + data={'imageBase64': 'iVBORw0KGgoAAAANSUhEUgAAAAUA'}, # dummy + format='json', + ) + self.assertEqual(resp.status_code, status.HTTP_200_OK) + mock_search.assert_called() + + @patch('api_v1.views.tiia.search_image_url_in_tencent_tiia') + def test_search_image_both_params_url_wins(self, mock_search): + self.client.force_authenticate(self.user) + mock_search.return_value = {'Candidates': []} + + resp = self.client.post( + '/api/v1/tiia/search-image/', + data={ + 'imageUrl': 'https://example.com/a.png', + 'imageBase64': 'iVBORw0KGgoAAAANSUhEUgAAAAUA', + }, + format='json', + ) + self.assertEqual(resp.status_code, status.HTTP_200_OK) + _, kwargs = mock_search.call_args + self.assertEqual(kwargs.get('image_url'), 'https://example.com/a.png') def test_search_image_missing_image_url(self): self.client.force_authenticate(self.user) diff --git a/api_v1/views/tiia.py b/api_v1/views/tiia.py index 058747d..0715a7a 100644 --- a/api_v1/views/tiia.py +++ b/api_v1/views/tiia.py @@ -34,16 +34,16 @@ class TiiaSearchImageView(APIView): permission_classes = [IsAuthenticated] def post(self, request): - image_url = ( - request.data.get('imageUrl') - or request.data.get('ImageUrl') - or request.data.get('image_url') - or request.data.get('ImageURL') - ) + # Keep API params consistent with existing style for this endpoint: + # - imageUrl + # - imageBase64 + image_url = request.data.get('imageUrl') + image_base64 = request.data.get('imageBase64') try: limiter = SimpleRateLimiter(qps=float(getattr(settings, 'TENCENTCLOUD_TIIA_QPS', 10))) resp = search_image_url_in_tencent_tiia( - image_url=str(image_url or ''), + image_url=str(image_url or '') if image_url is not None else None, + image_base64=str(image_base64 or '') if image_base64 is not None else None, rate_limiter=limiter, ) return Response(resp, status=status.HTTP_200_OK)