forked from erp-dev/erp
fear: added imageBase64 param for image search api
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user