mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
feat(dashscope): add embeddings support via OpenAI-compatible endpoint
This commit is contained in:
parent
0bcff0214a
commit
2a0563ab7a
7 changed files with 379 additions and 0 deletions
|
|
@ -1873,6 +1873,9 @@ if TYPE_CHECKING:
|
|||
from .llms.dashscope.chat.transformation import (
|
||||
DashScopeChatConfig as DashScopeChatConfig,
|
||||
)
|
||||
from .llms.dashscope.embed.transformation import (
|
||||
DashScopeEmbeddingConfig as DashScopeEmbeddingConfig,
|
||||
)
|
||||
from .llms.moonshot.chat.transformation import (
|
||||
MoonshotChatConfig as MoonshotChatConfig,
|
||||
)
|
||||
|
|
|
|||
28
litellm/llms/dashscope/common_utils.py
Normal file
28
litellm/llms/dashscope/common_utils.py
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
"""
|
||||
Common utilities for the DashScope LLM provider.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
|
||||
class DashScopeError(BaseLLMException):
|
||||
"""Exception class for DashScope provider errors."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: Optional[httpx.Headers] = None,
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.headers = headers or httpx.Headers()
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
headers=dict(self.headers),
|
||||
)
|
||||
7
litellm/llms/dashscope/embed/__init__.py
Normal file
7
litellm/llms/dashscope/embed/__init__.py
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
"""
|
||||
DashScope Embedding Module
|
||||
"""
|
||||
|
||||
from .transformation import DashScopeEmbeddingConfig
|
||||
|
||||
__all__ = ["DashScopeEmbeddingConfig"]
|
||||
190
litellm/llms/dashscope/embed/transformation.py
Normal file
190
litellm/llms/dashscope/embed/transformation.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
"""
|
||||
Transformation logic from OpenAI /v1/embeddings format to DashScope's /v1/embeddings format.
|
||||
|
||||
Supports
|
||||
- text-embedding-v4
|
||||
- text-embedding-v3
|
||||
|
||||
Endpoint
|
||||
- https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings
|
||||
|
||||
Docs - https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
"""
|
||||
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
|
||||
from ..common_utils import DashScopeError
|
||||
|
||||
DEFAULT_API_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1"
|
||||
|
||||
|
||||
class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
|
||||
"""
|
||||
Reference: https://help.aliyun.com/zh/model-studio/text-embedding-synchronous-api
|
||||
|
||||
DashScope exposes an OpenAI-compatible /v1/embeddings endpoint, so the
|
||||
request and response shapes are nearly identical to OpenAI's.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> List[str]:
|
||||
# DashScope's compatible-mode embeddings API accepts the same params as OpenAI.
|
||||
# `dimensions` / `encoding_format` are only honored by text-embedding-v3 / v4;
|
||||
# earlier versions silently ignore them server-side.
|
||||
return ["dimensions", "encoding_format"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool = False,
|
||||
) -> dict:
|
||||
supported = self.get_supported_openai_params(model)
|
||||
for k, v in non_default_params.items():
|
||||
if v is None:
|
||||
continue
|
||||
if k in supported:
|
||||
optional_params[k] = v
|
||||
elif not drop_params:
|
||||
raise ValueError(f"Unsupported parameter for DashScope embeddings: {k}")
|
||||
return optional_params
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> dict:
|
||||
resolved_key = api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
if not resolved_key:
|
||||
raise ValueError(
|
||||
"Missing API key for DashScope. Set DASHSCOPE_API_KEY environment "
|
||||
"variable or pass api_key parameter."
|
||||
)
|
||||
default_headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {resolved_key}",
|
||||
}
|
||||
return {**default_headers, **headers}
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
stream: Optional[bool] = None,
|
||||
) -> str:
|
||||
base = api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
|
||||
base = base.rstrip("/")
|
||||
if base.endswith("/embeddings"):
|
||||
return base
|
||||
return f"{base}/embeddings"
|
||||
|
||||
def transform_embedding_request(
|
||||
self,
|
||||
model: str,
|
||||
input: AllEmbeddingInputValues,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
data: dict = {
|
||||
"model": model,
|
||||
"input": input,
|
||||
}
|
||||
for key in ("dimensions", "encoding_format"):
|
||||
value = optional_params.get(key)
|
||||
if value is not None:
|
||||
data[key] = value
|
||||
return data
|
||||
|
||||
def transform_embedding_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: EmbeddingResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: Optional[str],
|
||||
request_data: dict,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
response_json = raw_response.json()
|
||||
except Exception as e:
|
||||
raise DashScopeError(
|
||||
status_code=raw_response.status_code,
|
||||
message=f"Failed to parse DashScope response as JSON: {str(e)}",
|
||||
)
|
||||
|
||||
logging_obj.post_call(
|
||||
input=request_data.get("input"),
|
||||
api_key=api_key,
|
||||
additional_args={"complete_input_dict": request_data},
|
||||
original_response=response_json,
|
||||
)
|
||||
|
||||
if "error" in response_json:
|
||||
error = response_json["error"]
|
||||
message = (
|
||||
error.get("message", str(error))
|
||||
if isinstance(error, dict)
|
||||
else str(error)
|
||||
)
|
||||
raise DashScopeError(
|
||||
status_code=raw_response.status_code,
|
||||
message=message,
|
||||
)
|
||||
|
||||
model_response.object = "list"
|
||||
model_response.data = response_json.get("data", [])
|
||||
model_response.model = response_json.get("model", model)
|
||||
|
||||
usage = response_json.get("usage") or {}
|
||||
prompt_tokens = usage.get("prompt_tokens", 0)
|
||||
total_tokens = usage.get("total_tokens", prompt_tokens)
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=total_tokens,
|
||||
),
|
||||
)
|
||||
|
||||
if "id" in response_json:
|
||||
setattr(model_response, "id", response_json["id"])
|
||||
|
||||
return model_response
|
||||
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: Union[dict, httpx.Headers],
|
||||
) -> BaseLLMException:
|
||||
if isinstance(headers, dict):
|
||||
headers = httpx.Headers(headers)
|
||||
return DashScopeError(
|
||||
status_code=status_code,
|
||||
message=error_message,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
@ -5645,6 +5645,33 @@ def embedding( # noqa: PLR0915
|
|||
aembedding=aembedding,
|
||||
headers=headers,
|
||||
)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
dashscope_key = (
|
||||
api_key or litellm.api_key or get_secret_str("DASHSCOPE_API_KEY")
|
||||
)
|
||||
if dashscope_key is None:
|
||||
raise ValueError(
|
||||
"Missing API key for DashScope. Set DASHSCOPE_API_KEY environment variable or pass api_key parameter."
|
||||
)
|
||||
if extra_headers is not None and isinstance(extra_headers, dict):
|
||||
headers = extra_headers
|
||||
else:
|
||||
headers = {}
|
||||
response = base_llm_http_handler.embedding(
|
||||
model=model,
|
||||
input=input,
|
||||
timeout=timeout,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
logging_obj=logging,
|
||||
api_base=api_base,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
model_response=EmbeddingResponse(),
|
||||
api_key=dashscope_key,
|
||||
client=client,
|
||||
aembedding=aembedding,
|
||||
headers=headers,
|
||||
)
|
||||
elif custom_llm_provider == "ovhcloud":
|
||||
api_key = api_key or litellm.api_key or get_secret_str("OVHCLOUD_API_KEY")
|
||||
api_base = (
|
||||
|
|
|
|||
|
|
@ -8368,6 +8368,12 @@ class ProviderConfigManager:
|
|||
)
|
||||
|
||||
return VolcEngineEmbeddingConfig()
|
||||
elif litellm.LlmProviders.DASHSCOPE == provider:
|
||||
from litellm.llms.dashscope.embed.transformation import (
|
||||
DashScopeEmbeddingConfig,
|
||||
)
|
||||
|
||||
return DashScopeEmbeddingConfig()
|
||||
elif litellm.LlmProviders.OVHCLOUD == provider:
|
||||
return litellm.OVHCloudEmbeddingConfig()
|
||||
elif litellm.LlmProviders.SNOWFLAKE == provider:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,118 @@
|
|||
"""
|
||||
Unit tests for DashScope embedding transformation.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.llms.dashscope.common_utils import DashScopeError
|
||||
from litellm.llms.dashscope.embed.transformation import (
|
||||
DEFAULT_API_BASE,
|
||||
DashScopeEmbeddingConfig,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
def test_validate_environment_and_url():
|
||||
config = DashScopeEmbeddingConfig()
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="text-embedding-v4",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="sk-test",
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer sk-test"
|
||||
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key="sk-test",
|
||||
model="text-embedding-v4",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == f"{DEFAULT_API_BASE}/embeddings"
|
||||
|
||||
|
||||
def test_transform_embedding_request():
|
||||
config = DashScopeEmbeddingConfig()
|
||||
data = config.transform_embedding_request(
|
||||
model="text-embedding-v4",
|
||||
input=["风急天高猿啸哀"],
|
||||
optional_params={"dimensions": 1024, "encoding_format": "float"},
|
||||
headers={},
|
||||
)
|
||||
assert data == {
|
||||
"model": "text-embedding-v4",
|
||||
"input": ["风急天高猿啸哀"],
|
||||
"dimensions": 1024,
|
||||
"encoding_format": "float",
|
||||
}
|
||||
|
||||
|
||||
def test_transform_embedding_response_success():
|
||||
config = DashScopeEmbeddingConfig()
|
||||
payload = {
|
||||
"data": [
|
||||
{"embedding": [0.1, 0.2], "index": 0, "object": "embedding"},
|
||||
],
|
||||
"model": "text-embedding-v4",
|
||||
"object": "list",
|
||||
"usage": {"prompt_tokens": 5, "total_tokens": 5},
|
||||
"id": "73591b79-xxxx",
|
||||
}
|
||||
raw = httpx.Response(
|
||||
status_code=200,
|
||||
content=json.dumps(payload).encode("utf-8"),
|
||||
request=httpx.Request("POST", "https://example.com"),
|
||||
)
|
||||
result = config.transform_embedding_response(
|
||||
model="text-embedding-v4",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
api_key="sk-x",
|
||||
request_data={"input": ["a"]},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert result.model == "text-embedding-v4"
|
||||
assert len(result.data) == 1
|
||||
assert result.usage.prompt_tokens == 5
|
||||
|
||||
|
||||
def test_transform_embedding_response_error():
|
||||
config = DashScopeEmbeddingConfig()
|
||||
payload = {
|
||||
"error": {
|
||||
"message": "Incorrect API key provided.",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
}
|
||||
raw = httpx.Response(
|
||||
status_code=401,
|
||||
content=json.dumps(payload).encode("utf-8"),
|
||||
request=httpx.Request("POST", "https://example.com"),
|
||||
)
|
||||
with pytest.raises(DashScopeError) as exc:
|
||||
config.transform_embedding_response(
|
||||
model="text-embedding-v4",
|
||||
raw_response=raw,
|
||||
model_response=EmbeddingResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
api_key="sk-bad",
|
||||
request_data={"input": ["a"]},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert exc.value.status_code == 401
|
||||
assert "Incorrect API key" in exc.value.message
|
||||
Loading…
Add table
Reference in a new issue