-
Notifications
You must be signed in to change notification settings - Fork 8.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: support siliconflow rerank (#7337)
- Loading branch information
Showing
6 changed files
with
146 additions
and
1 deletion.
There are no files selected for viewing
Empty file.
4 changes: 4 additions & 0 deletions
4
api/core/model_runtime/model_providers/siliconflow/rerank/bce-reranker-base_v1.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
model: netease-youdao/bce-reranker-base_v1 | ||
model_type: rerank | ||
model_properties: | ||
context_size: 512 |
4 changes: 4 additions & 0 deletions
4
api/core/model_runtime/model_providers/siliconflow/rerank/bge-reranker-v2-m3.yaml
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,4 @@ | ||
model: BAAI/bge-reranker-v2-m3 | ||
model_type: rerank | ||
model_properties: | ||
context_size: 8192 |
87 changes: 87 additions & 0 deletions
87
api/core/model_runtime/model_providers/siliconflow/rerank/rerank.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,87 @@ | ||
from typing import Optional | ||
|
||
import httpx | ||
|
||
from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult | ||
from core.model_runtime.errors.invoke import ( | ||
InvokeAuthorizationError, | ||
InvokeBadRequestError, | ||
InvokeConnectionError, | ||
InvokeError, | ||
InvokeRateLimitError, | ||
InvokeServerUnavailableError, | ||
) | ||
from core.model_runtime.errors.validate import CredentialsValidateFailedError | ||
from core.model_runtime.model_providers.__base.rerank_model import RerankModel | ||
|
||
|
||
class SiliconflowRerankModel(RerankModel): | ||
|
||
def _invoke(self, model: str, credentials: dict, query: str, docs: list[str], | ||
score_threshold: Optional[float] = None, top_n: Optional[int] = None, | ||
user: Optional[str] = None) -> RerankResult: | ||
if len(docs) == 0: | ||
return RerankResult(model=model, docs=[]) | ||
|
||
base_url = credentials.get('base_url', 'https://api.siliconflow.cn/v1') | ||
if base_url.endswith('/'): | ||
base_url = base_url[:-1] | ||
try: | ||
response = httpx.post( | ||
base_url + '/rerank', | ||
json={ | ||
"model": model, | ||
"query": query, | ||
"documents": docs, | ||
"top_n": top_n, | ||
"return_documents": True | ||
}, | ||
headers={"Authorization": f"Bearer {credentials.get('api_key')}"} | ||
) | ||
response.raise_for_status() | ||
results = response.json() | ||
|
||
rerank_documents = [] | ||
for result in results['results']: | ||
rerank_document = RerankDocument( | ||
index=result['index'], | ||
text=result['document']['text'], | ||
score=result['relevance_score'], | ||
) | ||
if score_threshold is None or result['relevance_score'] >= score_threshold: | ||
rerank_documents.append(rerank_document) | ||
|
||
return RerankResult(model=model, docs=rerank_documents) | ||
except httpx.HTTPStatusError as e: | ||
raise InvokeServerUnavailableError(str(e)) | ||
|
||
def validate_credentials(self, model: str, credentials: dict) -> None: | ||
try: | ||
|
||
self._invoke( | ||
model=model, | ||
credentials=credentials, | ||
query="What is the capital of the United States?", | ||
docs=[ | ||
"Carson City is the capital city of the American state of Nevada. At the 2010 United States " | ||
"Census, Carson City had a population of 55,274.", | ||
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that " | ||
"are a political division controlled by the United States. Its capital is Saipan.", | ||
], | ||
score_threshold=0.8 | ||
) | ||
except Exception as ex: | ||
raise CredentialsValidateFailedError(str(ex)) | ||
|
||
@property | ||
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]: | ||
""" | ||
Map model invoke error to unified error | ||
""" | ||
return { | ||
InvokeConnectionError: [httpx.ConnectError], | ||
InvokeServerUnavailableError: [httpx.RemoteProtocolError], | ||
InvokeRateLimitError: [], | ||
InvokeAuthorizationError: [httpx.HTTPStatusError], | ||
InvokeBadRequestError: [httpx.RequestError] | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
49 changes: 49 additions & 0 deletions
49
api/tests/integration_tests/model_runtime/siliconflow/test_rerank.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,49 @@ | ||
import os | ||
|
||
import pytest | ||
|
||
from core.model_runtime.entities.rerank_entities import RerankResult | ||
from core.model_runtime.errors.validate import CredentialsValidateFailedError | ||
from core.model_runtime.model_providers.siliconflow.rerank.rerank import SiliconflowRerankModel | ||
|
||
|
||
def test_validate_credentials(): | ||
model = SiliconflowRerankModel() | ||
|
||
with pytest.raises(CredentialsValidateFailedError): | ||
model.validate_credentials( | ||
model="BAAI/bge-reranker-v2-m3", | ||
credentials={ | ||
"api_key": "invalid_key" | ||
}, | ||
) | ||
|
||
model.validate_credentials( | ||
model="BAAI/bge-reranker-v2-m3", | ||
credentials={ | ||
"api_key": os.environ.get("API_KEY"), | ||
}, | ||
) | ||
|
||
|
||
def test_invoke_model(): | ||
model = SiliconflowRerankModel() | ||
|
||
result = model.invoke( | ||
model='BAAI/bge-reranker-v2-m3', | ||
credentials={ | ||
"api_key": os.environ.get("API_KEY"), | ||
}, | ||
query="Who is Kasumi?", | ||
docs=[ | ||
"Kasumi is a girl's name of Japanese origin meaning \"mist\".", | ||
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ", | ||
"and she leads a team named PopiParty." | ||
], | ||
score_threshold=0.8 | ||
) | ||
|
||
assert isinstance(result, RerankResult) | ||
assert len(result.docs) == 1 | ||
assert result.docs[0].index == 0 | ||
assert result.docs[0].score >= 0.8 |