diff --git a/dashscope/__init__.py b/dashscope/__init__.py index 637d5de..92cd259 100644 --- a/dashscope/__init__.py +++ b/dashscope/__init__.py @@ -46,7 +46,7 @@ from dashscope.files import Files from dashscope.models import Models from dashscope.nlp.understanding import Understanding -from dashscope.rerank.text_rerank import TextReRank +from dashscope.rerank import AioTextReRank, TextReRank from dashscope.threads import ( MessageFile, Messages, @@ -106,6 +106,7 @@ "list_tokenizers", "Application", "TextReRank", + "AioTextReRank", "Assistants", "Threads", "Messages", diff --git a/dashscope/aigc/image_synthesis.py b/dashscope/aigc/image_synthesis.py index 41bd060..0491d7e 100644 --- a/dashscope/aigc/image_synthesis.py +++ b/dashscope/aigc/image_synthesis.py @@ -387,6 +387,7 @@ def wait( # type: ignore[override] task: Union[str, ImageSynthesisResponse], api_key: str = None, workspace: str = None, + **kwargs, ) -> ImageSynthesisResponse: """Wait for image(s) synthesis task to complete, and return the result. @@ -399,7 +400,12 @@ def wait( # type: ignore[override] Returns: ImageSynthesisResponse: The task result. """ - response = super().wait(task, api_key, workspace=workspace) + response = super().wait( + task, + api_key, + workspace=workspace, + **kwargs, + ) return ImageSynthesisResponse.from_api_response(response) @classmethod diff --git a/dashscope/aigc/video_synthesis.py b/dashscope/aigc/video_synthesis.py index 0bc5032..b976799 100644 --- a/dashscope/aigc/video_synthesis.py +++ b/dashscope/aigc/video_synthesis.py @@ -509,6 +509,7 @@ def wait( # type: ignore[override] task: Union[str, VideoSynthesisResponse], api_key: str = None, workspace: str = None, + **kwargs, ) -> VideoSynthesisResponse: """Wait for video synthesis task to complete, and return the result. @@ -521,7 +522,12 @@ def wait( # type: ignore[override] Returns: VideoSynthesisResponse: The task result. """ - response = super().wait(task, api_key, workspace=workspace) + response = super().wait( + task, + api_key, + workspace=workspace, + **kwargs, + ) return VideoSynthesisResponse.from_api_response(response) @classmethod diff --git a/dashscope/api_entities/dashscope_response.py b/dashscope/api_entities/dashscope_response.py index 71f16b8..58ade6b 100644 --- a/dashscope/api_entities/dashscope_response.py +++ b/dashscope/api_entities/dashscope_response.py @@ -59,7 +59,12 @@ def setattr(self, attr, value): return super().__setitem__(attr, value) def __getattr__(self, attr): - return self[attr] + try: + return self[attr] + except KeyError: + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {attr!r}", + ) from None def __setattr__(self, attr, value): self[attr] = value diff --git a/dashscope/embeddings/batch_text_embedding.py b/dashscope/embeddings/batch_text_embedding.py index d3af2d1..c2c9b9c 100644 --- a/dashscope/embeddings/batch_text_embedding.py +++ b/dashscope/embeddings/batch_text_embedding.py @@ -143,6 +143,7 @@ def wait( # type: ignore[override] task: Union[str, BatchTextEmbeddingResponse], api_key: str = None, workspace: str = None, + **kwargs, ) -> BatchTextEmbeddingResponse: """Wait for async text embedding task to complete, and return the result. # noqa: E501 @@ -155,7 +156,12 @@ def wait( # type: ignore[override] Returns: AsyncTextEmbeddingResponse: The task result. """ - response = super().wait(task, api_key, workspace=workspace) + response = super().wait( + task, + api_key, + workspace=workspace, + **kwargs, + ) return BatchTextEmbeddingResponse.from_api_response(response) @classmethod diff --git a/dashscope/rerank/__init__.py b/dashscope/rerank/__init__.py index e69de29..2f0829c 100644 --- a/dashscope/rerank/__init__.py +++ b/dashscope/rerank/__init__.py @@ -0,0 +1,6 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +from dashscope.rerank.text_rerank import AioTextReRank, TextReRank + +__all__ = ["AioTextReRank", "TextReRank"] diff --git a/dashscope/rerank/text_rerank.py b/dashscope/rerank/text_rerank.py index b152ae6..3a4b31e 100644 --- a/dashscope/rerank/text_rerank.py +++ b/dashscope/rerank/text_rerank.py @@ -1,13 +1,46 @@ # -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. -from typing import List +from typing import Any, Dict, List, Tuple from dashscope.api_entities.dashscope_response import ReRankResponse -from dashscope.client.base_api import BaseApi +from dashscope.client.base_api import BaseAioApi, BaseApi from dashscope.common.error import InputRequired, ModelRequired from dashscope.common.utils import _get_task_group_and_task +__all__ = ["TextReRank", "AioTextReRank"] + + +def _build_rerank_request( + model: str, + query: str, + documents: List[str], + return_documents: bool = None, + top_n: int = None, + instruct: str = None, + **kwargs, +) -> Tuple[str, str, Dict[str, Any], Dict[str, Any]]: + if query is None or documents is None or not documents: + raise InputRequired("query and documents are required!") + if model is None or not model: + raise ModelRequired("Model is required!") + + task_group, function = _get_task_group_and_task(__name__) + rerank_input = { + "query": query, + "documents": documents, + } + parameters = {} + if return_documents is not None: + parameters["return_documents"] = return_documents + if top_n is not None: + parameters["top_n"] = top_n + if instruct is not None: + parameters["instruct"] = instruct + parameters = {**parameters, **kwargs} + + return task_group, function, rerank_input, parameters + class TextReRank(BaseApi): task = "text-rerank" @@ -41,8 +74,8 @@ def call( # type: ignore[override] # pylint: disable=arguments-renamed documents (List[str]): The documents to rank. return_documents(bool, `optional`): enable return origin documents, system default is false. - top_n(int, `optional`): how many documents to return, default return # noqa: E501 - all the documents. + top_n(int, `optional`): how many documents to return, + default return all the documents. api_key (str, optional): The DashScope api key. Defaults to None. instruct (str, optional): Custom task instruction to guide ranking strategy. English recommended. @@ -55,23 +88,15 @@ def call( # type: ignore[override] # pylint: disable=arguments-renamed RerankResponse: The rerank result. """ - if query is None or documents is None or not documents: - raise InputRequired("query and documents are required!") - if model is None or not model: - raise ModelRequired("Model is required!") - task_group, function = _get_task_group_and_task(__name__) - input = { # pylint: disable=redefined-builtin - "query": query, - "documents": documents, - } - parameters = {} - if return_documents is not None: - parameters["return_documents"] = return_documents - if top_n is not None: - parameters["top_n"] = top_n - if instruct is not None: - parameters["instruct"] = instruct - parameters = {**parameters, **kwargs} + task_group, function, rerank_input, parameters = _build_rerank_request( + model=model, + query=query, + documents=documents, + return_documents=return_documents, + top_n=top_n, + instruct=instruct, + **kwargs, + ) response = super().call( model=model, @@ -79,7 +104,73 @@ def call( # type: ignore[override] # pylint: disable=arguments-renamed task=TextReRank.task, function=function, api_key=api_key, - input=input, + input=rerank_input, + **parameters, # type: ignore[arg-type] + ) + + return ReRankResponse.from_api_response(response) + + +class AioTextReRank(BaseAioApi): + task = "text-rerank" + """Async API for rerank models.""" + + Models = TextReRank.Models + + @classmethod + # pylint: disable=arguments-renamed + async def call( # type: ignore[override] + cls, + model: str, + query: str, + documents: List[str], + return_documents: bool = None, + top_n: int = None, + api_key: str = None, + workspace: str = None, + instruct: str = None, + **kwargs, + ) -> ReRankResponse: + """Calling rerank service asynchronously. + + Args: + model (str): The model to use. + query (str): The query string. + documents (List[str]): The documents to rank. + return_documents(bool, `optional`): enable return origin documents, + system default is false. + top_n(int, `optional`): how many documents to return, + default return all the documents. + api_key (str, optional): The DashScope api key. Defaults to None. + workspace (str, optional): The DashScope workspace id. + instruct (str, optional): Custom task instruction to guide + ranking strategy. English recommended. + + Raises: + InputRequired: The query and documents are required. + ModelRequired: The model is required. + + Returns: + RerankResponse: The rerank result. + """ + task_group, function, rerank_input, parameters = _build_rerank_request( + model=model, + query=query, + documents=documents, + return_documents=return_documents, + top_n=top_n, + instruct=instruct, + **kwargs, + ) + + response = await super().call( + model=model, + task_group=task_group, + task=AioTextReRank.task, + function=function, + api_key=api_key, + workspace=workspace, + input=rerank_input, **parameters, # type: ignore[arg-type] ) diff --git a/dashscope/utils/oss_utils.py b/dashscope/utils/oss_utils.py index 216333e..6b5e476 100644 --- a/dashscope/utils/oss_utils.py +++ b/dashscope/utils/oss_utils.py @@ -129,6 +129,23 @@ def get_upload_certificate( return super().get(None, api_key, params=params, **kwargs) # type: ignore[return-value] # pylint: disable=line-too-long # noqa: E501 +def _resolve_file_uri_path(file_uri: str): + parse_result = urlparse(file_uri) + if parse_result.netloc: + file_path = parse_result.netloc + unquote_plus(parse_result.path) + else: + file_path = unquote_plus(parse_result.path) + + if ( + file_path.startswith("/") + and len(file_path) > 2 + and file_path[2] == ":" + ): + file_path = file_path[1:] + + return os.path.expanduser(file_path) + + def upload_file( model: str, upload_path: str, @@ -136,11 +153,7 @@ def upload_file( upload_certificate: dict = None, ): if upload_path.startswith(FILE_PATH_SCHEMA): - parse_result = urlparse(upload_path) - if parse_result.netloc: - file_path = parse_result.netloc + unquote_plus(parse_result.path) - else: - file_path = unquote_plus(parse_result.path) + file_path = _resolve_file_uri_path(upload_path) if os.path.exists(file_path): file_url, _ = OssUtils.upload( model=model, @@ -184,11 +197,7 @@ def check_and_upload_local( is the certificate (newly obtained or passed in) """ if content.startswith(FILE_PATH_SCHEMA): - parse_result = urlparse(content) - if parse_result.netloc: - file_path = parse_result.netloc + unquote_plus(parse_result.path) - else: - file_path = unquote_plus(parse_result.path) + file_path = _resolve_file_uri_path(content) if os.path.isfile(file_path): file_url, cert = OssUtils.upload( model=model, @@ -201,9 +210,10 @@ def check_and_upload_local( f"Uploading file: {content} failed", ) return True, file_url, cert - elif content.startswith("oss://"): + raise InvalidInput(f"The file: {file_path} is not exists!") + if content.startswith("oss://"): return True, content, upload_certificate - elif not content.startswith("http"): + if not content.startswith("http"): content = os.path.expanduser(content) if os.path.isfile(content): file_url, cert = OssUtils.upload( diff --git a/tests/unit/test_dashscope_response.py b/tests/unit/test_dashscope_response.py new file mode 100644 index 0000000..ed7a135 --- /dev/null +++ b/tests/unit/test_dashscope_response.py @@ -0,0 +1,21 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +from dashscope.api_entities.dashscope_response import DictMixin + + +class TestDictMixin: + def test_getattr_missing_key_raises_attribute_error(self): + response = DictMixin(existing="value") + + try: + response.missing + except AttributeError: + return + + raise AssertionError("Missing attribute should raise AttributeError") + + def test_getattr_existing_key_returns_value(self): + response = DictMixin(existing="value") + + assert response.existing == "value" diff --git a/tests/unit/test_oss_utils.py b/tests/unit/test_oss_utils.py new file mode 100644 index 0000000..505d384 --- /dev/null +++ b/tests/unit/test_oss_utils.py @@ -0,0 +1,147 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. + +from http import HTTPStatus + +import pytest + +from dashscope.common.error import InvalidInput +from dashscope.utils import oss_utils +from dashscope.utils.oss_utils import OssUtils + + +class FakeUploadResponse: + status_code = HTTPStatus.OK + headers = {} + + +class FakeSession: + captured_file = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def post(self, url, files, data, headers, timeout): + assert url == "https://oss.example.com" + assert data["key"] == "test-dir/dogs.jpg" + assert headers["Accept"] == "application/json" + assert timeout == 3600 + + FakeSession.captured_file = files["file"] + assert not FakeSession.captured_file.closed + return FakeUploadResponse() + + +class TestOssUtils: + def test_upload_closes_opened_file(self, monkeypatch): + upload_certificate = { + "oss_access_key_id": "access-key-id", + "signature": "signature", + "policy": "policy", + "upload_dir": "test-dir", + "x_oss_object_acl": "private", + "x_oss_forbid_overwrite": "true", + "upload_host": "https://oss.example.com", + } + FakeSession.captured_file = None + monkeypatch.setattr(oss_utils.requests, "Session", FakeSession) + + file_url, returned_certificate = OssUtils.upload( + model="test-model", + file_path="tests/data/dogs.jpg", + api_key="test-api-key", + upload_certificate=upload_certificate, + ) + + assert file_url == "oss://test-dir/dogs.jpg" + assert returned_certificate is upload_certificate + assert FakeSession.captured_file is not None + assert FakeSession.captured_file.closed + + def test_check_and_upload_local_uploads_relative_file_uri( + self, + monkeypatch, + ): + captured_file_path = {} + + def fake_isfile(file_path): + captured_file_path["value"] = file_path + return True + + def fake_upload(model, file_path, api_key, upload_certificate): + assert model == "test-model" + assert api_key == "test-api-key" + assert upload_certificate == {"cert": "value"} + assert file_path == "test_video_frames/frame_0000.jpg" + return "oss://test-dir/frame_0000.jpg", {"cert": "value"} + + monkeypatch.setattr(oss_utils.os.path, "isfile", fake_isfile) + monkeypatch.setattr(OssUtils, "upload", fake_upload) + + is_upload, file_url, certificate = oss_utils.check_and_upload_local( + model="test-model", + content="file://test_video_frames/frame_0000.jpg", + api_key="test-api-key", + upload_certificate={"cert": "value"}, + ) + + assert is_upload + assert file_url == "oss://test-dir/frame_0000.jpg" + assert certificate == {"cert": "value"} + assert ( + captured_file_path["value"] == "test_video_frames/frame_0000.jpg" + ) + + def test_check_and_upload_local_supports_windows_absolute_file_uri( + self, + monkeypatch, + ): + captured_file_path = {} + + def fake_isfile(file_path): + captured_file_path["value"] = file_path + return True + + def fake_upload( + model, + file_path, + api_key, + upload_certificate, + ): + assert model == "test-model" + assert file_path == "C:/Users/test/frame_0000.jpg" + assert api_key == "test-api-key" + return "oss://test-dir/frame_0000.jpg", upload_certificate + + monkeypatch.setattr(oss_utils.os.path, "isfile", fake_isfile) + monkeypatch.setattr(OssUtils, "upload", fake_upload) + + is_upload, file_url, _ = oss_utils.check_and_upload_local( + model="test-model", + content="file:///C:/Users/test/frame_0000.jpg", + api_key="test-api-key", + ) + + assert is_upload + assert file_url == "oss://test-dir/frame_0000.jpg" + assert captured_file_path["value"] == "C:/Users/test/frame_0000.jpg" + + def test_check_and_upload_local_raises_when_file_uri_not_found( + self, + monkeypatch, + ): + monkeypatch.setattr( + oss_utils.os.path, + "isfile", + lambda file_path: False, + ) + + with pytest.raises(InvalidInput): + oss_utils.check_and_upload_local( + model="test-model", + content="file://missing/frame_0000.jpg", + api_key="test-api-key", + ) diff --git a/tests/unit/test_rerank.py b/tests/unit/test_rerank.py index d2afd34..b742794 100644 --- a/tests/unit/test_rerank.py +++ b/tests/unit/test_rerank.py @@ -1,10 +1,11 @@ # -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. +import asyncio import json import uuid -from dashscope import TextReRank +from dashscope import AioTextReRank, TextReRank from tests.unit.mock_request_base import MockServerBase from tests.unit.mock_server import MockServer @@ -62,3 +63,61 @@ def test_call(self, mock_server: MockServer): assert len(response.output["results"]) == 2 assert response.output["results"][0]["index"] == 1 assert response.output["results"][1]["document"]["text"] == "黑龙江离俄罗斯很近" + + def test_aio_call(self, mock_server: MockServer): + response_body = { + "output": { + "results": [ + { + "index": 1, + "relevance_score": 0.987654, + "document": { + "text": "哈尔滨是中国黑龙江省的省会，位于中国东北", + }, + }, + { + "index": 0, + "relevance_score": 0.876543, + "document": { + "text": "黑龙江离俄罗斯很近", + }, + }, + ], + }, + "usage": { + "input_tokens": 1279, + }, + "request_id": "b042e72d-7994-97dd-b3d2-7ee7e0140525", + } + mock_server.responses.put(json.dumps(response_body)) + model = str(uuid.uuid4()) + query = str(uuid.uuid4()) + documents = [ + str(uuid.uuid4()), + str(uuid.uuid4()), + str(uuid.uuid4()), + str(uuid.uuid4()), + ] + + response = asyncio.run( + AioTextReRank.call( + model=model, + query=query, + documents=documents, + return_documents=True, + top_n=2, + instruct="Rank the documents by relevance.", + ), + ) + + req = mock_server.requests.get(block=True) + assert req["path"] == "/api/v1/services/rerank/text-rerank/text-rerank" + assert req["body"]["parameters"] == { + "return_documents": True, + "top_n": 2, + "instruct": "Rank the documents by relevance.", + } + assert req["body"]["input"] == {"query": query, "documents": documents} + assert response.usage["input_tokens"] == 1279 + assert len(response.output["results"]) == 2 + assert response.output["results"][0]["index"] == 1