mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 04:21:34 +00:00
feat(perplexity): add embedding support for pplx-embed-v1 models (#22610)
* feat: add Perplexity embedding support (pplx-embed-v1) Add support for Perplexity AI's embedding models via the LLM HTTP handler: Models: - pplx-embed-v1-0.6b (1024 dims, 32K context, $0.004/1M tokens) - pplx-embed-v1-4b (2560 dims, 32K context, $0.03/1M tokens) Implementation: - PerplexityEmbeddingConfig in litellm/llms/perplexity/embedding/ - Registered in ProviderConfigManager, __init__.py lazy imports, main.py dispatch - Model pricing added to model_prices_and_context_window.json - Supports dimensions and encoding_format parameters - Uses base_llm_http_handler.embedding() pattern Tests: - 19 unit tests covering transformation, params, URLs, provider config, model info Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * docs: add Perplexity AI embeddings documentation - Create providers/perplexity_embedding.md with SDK and proxy usage examples - Convert Perplexity from flat doc to category in sidebars.js - Category includes existing chat/responses doc + new embeddings doc - Covers pplx-embed-v1-0.6b and pplx-embed-v1-4b models - Documents supported parameters (dimensions, encoding_format) - Includes proxy config and curl examples Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> * fix: decode Perplexity base64_int8 embeddings to OpenAI-format float arrays Perplexity returns embeddings as base64-encoded signed int8 values by default, not float arrays like OpenAI. This commit adds decoding in transform_embedding_response so the proxy returns standard OpenAI-compatible float arrays (normalized to [-1, 1]). - Added _decode_base64_embedding() static method - Handles both base64 strings (decoded) and float lists (passthrough) - Added 3 new tests for base64 decoding + passthrough Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
co-authored by
Cursor Agent
Ishaan Jaff
Cursor Agent
Ishaan Jaff
parent
b8befb3403
commit
bfceb7fc3f
@@ -0,0 +1,134 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Perplexity Embeddings
|
||||
|
||||
https://docs.perplexity.ai/docs/embeddings/quickstart
|
||||
|
||||
LiteLLM supports Perplexity's pplx-embed embedding models for web-scale text retrieval.
|
||||
|
||||
## API Key
|
||||
|
||||
```python
|
||||
# env variable
|
||||
os.environ['PERPLEXITYAI_API_KEY']
|
||||
```
|
||||
|
||||
## Sample Usage - Embedding
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="SDK">
|
||||
|
||||
```python
|
||||
from litellm import embedding
|
||||
import os
|
||||
|
||||
os.environ['PERPLEXITYAI_API_KEY'] = ""
|
||||
|
||||
response = embedding(
|
||||
model="perplexity/pplx-embed-v1-0.6b",
|
||||
input=["good morning from litellm"],
|
||||
)
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="Proxy">
|
||||
|
||||
1. Setup config.yaml
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: pplx-embed-v1-0.6b
|
||||
litellm_params:
|
||||
model: perplexity/pplx-embed-v1-0.6b
|
||||
api_key: os.environ/PERPLEXITYAI_API_KEY
|
||||
- model_name: pplx-embed-v1-4b
|
||||
litellm_params:
|
||||
model: perplexity/pplx-embed-v1-4b
|
||||
api_key: os.environ/PERPLEXITYAI_API_KEY
|
||||
```
|
||||
|
||||
2. Start proxy
|
||||
|
||||
```bash
|
||||
litellm --config /path/to/config.yaml
|
||||
```
|
||||
|
||||
3. Test it!
|
||||
|
||||
```bash
|
||||
curl http://0.0.0.0:4000/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "pplx-embed-v1-0.6b",
|
||||
"input": ["good morning from litellm"]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Supported Parameters
|
||||
|
||||
Perplexity embeddings support the following optional parameters:
|
||||
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `dimensions` | int | Output embedding dimensions. 128–1024 for 0.6b models, 128–2560 for 4b models. Defaults to max. |
|
||||
| `encoding_format` | string | `"base64_int8"` (default) or `"base64_binary"` for compressed output. |
|
||||
|
||||
### Example with Parameters
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="SDK">
|
||||
|
||||
```python
|
||||
from litellm import embedding
|
||||
import os
|
||||
|
||||
os.environ['PERPLEXITYAI_API_KEY'] = ""
|
||||
|
||||
response = embedding(
|
||||
model="perplexity/pplx-embed-v1-4b",
|
||||
input=["Your text here"],
|
||||
dimensions=512,
|
||||
)
|
||||
print(f"Embedding dimensions: {len(response.data[0]['embedding'])}")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="Proxy">
|
||||
|
||||
```bash
|
||||
curl http://0.0.0.0:4000/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "pplx-embed-v1-4b",
|
||||
"input": ["Your text here"],
|
||||
"dimensions": 512
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Supported Models
|
||||
|
||||
All models listed on the [Perplexity Embeddings docs](https://docs.perplexity.ai/docs/embeddings/quickstart) are supported. Use `model=perplexity/<model-name>`.
|
||||
|
||||
| Model Name | Dimensions | Max Tokens | Price (per 1M tokens) | Function Call |
|
||||
|---|---|---|---|---|
|
||||
| pplx-embed-v1-0.6b | 1024 | 32K | $0.004 | `embedding(model="perplexity/pplx-embed-v1-0.6b", input)` |
|
||||
| pplx-embed-v1-4b | 2560 | 32K | $0.03 | `embedding(model="perplexity/pplx-embed-v1-4b", input)` |
|
||||
|
||||
### Key Specifications
|
||||
|
||||
- **Max texts per request:** 512
|
||||
- **Max tokens per input:** 32,768
|
||||
- **Combined request limit:** 120,000 tokens
|
||||
- **Matryoshka dimension reduction** — reduce dimensions to 128+ for faster search and reduced storage
|
||||
- **No instruction prefix required** — embed text directly
|
||||
- **Unnormalized embeddings** — use cosine similarity for comparison
|
||||
@@ -877,7 +877,14 @@ const sidebars = {
|
||||
"providers/openrouter",
|
||||
"providers/sarvam",
|
||||
"providers/ovhcloud",
|
||||
"providers/perplexity",
|
||||
{
|
||||
type: "category",
|
||||
label: "Perplexity AI",
|
||||
items: [
|
||||
"providers/perplexity",
|
||||
"providers/perplexity_embedding",
|
||||
]
|
||||
},
|
||||
"providers/petals",
|
||||
"providers/poe",
|
||||
"providers/publicai",
|
||||
|
||||
@@ -1429,6 +1429,7 @@ if TYPE_CHECKING:
|
||||
from .llms.voyage.embedding.transformation import VoyageEmbeddingConfig as VoyageEmbeddingConfig
|
||||
from .llms.voyage.embedding.transformation_contextual import VoyageContextualEmbeddingConfig as VoyageContextualEmbeddingConfig
|
||||
from .llms.infinity.embedding.transformation import InfinityEmbeddingConfig as InfinityEmbeddingConfig
|
||||
from .llms.perplexity.embedding.transformation import PerplexityEmbeddingConfig as PerplexityEmbeddingConfig
|
||||
from .llms.azure_ai.chat.transformation import AzureAIStudioConfig as AzureAIStudioConfig
|
||||
from .llms.mistral.chat.transformation import MistralConfig as MistralConfig
|
||||
from .llms.openai.responses.transformation import OpenAIResponsesAPIConfig as OpenAIResponsesAPIConfig
|
||||
|
||||
@@ -219,6 +219,7 @@ LLM_CONFIG_NAMES = (
|
||||
"VoyageEmbeddingConfig",
|
||||
"VoyageContextualEmbeddingConfig",
|
||||
"InfinityEmbeddingConfig",
|
||||
"PerplexityEmbeddingConfig",
|
||||
"AzureAIStudioConfig",
|
||||
"MistralConfig",
|
||||
"OpenAIResponsesAPIConfig",
|
||||
@@ -873,6 +874,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
||||
".llms.infinity.embedding.transformation",
|
||||
"InfinityEmbeddingConfig",
|
||||
),
|
||||
"PerplexityEmbeddingConfig": (
|
||||
".llms.perplexity.embedding.transformation",
|
||||
"PerplexityEmbeddingConfig",
|
||||
),
|
||||
"AzureAIStudioConfig": (
|
||||
".llms.azure_ai.chat.transformation",
|
||||
"AzureAIStudioConfig",
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Perplexity AI Embedding API
|
||||
|
||||
Docs: https://docs.perplexity.ai/api-reference/embeddings-post
|
||||
|
||||
Supports models:
|
||||
- pplx-embed-v1-0.6b (1024 dims, 32 K context)
|
||||
- pplx-embed-v1-4b (2560 dims, 32 K context)
|
||||
|
||||
Perplexity returns embeddings as base64-encoded signed int8 values by default.
|
||||
This module decodes them into float arrays for OpenAI-compatible responses.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import struct
|
||||
from typing import Any, Dict, 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
|
||||
|
||||
|
||||
class PerplexityEmbeddingError(BaseLLMException):
|
||||
def __init__(
|
||||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: Union[dict, httpx.Headers] = {},
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.request = httpx.Request(
|
||||
method="POST", url="https://api.perplexity.ai/v1/embeddings"
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request)
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
class PerplexityEmbeddingConfig(BaseEmbeddingConfig):
|
||||
"""
|
||||
Reference: https://docs.perplexity.ai/api-reference/embeddings-post
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
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:
|
||||
if api_base:
|
||||
if not api_base.endswith("/embeddings"):
|
||||
api_base = f"{api_base}/v1/embeddings"
|
||||
return api_base
|
||||
return "https://api.perplexity.ai/v1/embeddings"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
return [
|
||||
"dimensions",
|
||||
"encoding_format",
|
||||
]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
for k, v in non_default_params.items():
|
||||
if k == "dimensions":
|
||||
optional_params["dimensions"] = v
|
||||
elif k == "encoding_format":
|
||||
optional_params["encoding_format"] = v
|
||||
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:
|
||||
if api_key is None:
|
||||
api_key = get_secret_str("PERPLEXITYAI_API_KEY") or get_secret_str(
|
||||
"PERPLEXITY_API_KEY"
|
||||
)
|
||||
return {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def transform_embedding_request(
|
||||
self,
|
||||
model: str,
|
||||
input: AllEmbeddingInputValues,
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
return {
|
||||
"model": model,
|
||||
"input": input,
|
||||
**optional_params,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _decode_base64_embedding(embedding_value: Any) -> List[float]:
|
||||
"""
|
||||
Decode a Perplexity embedding into a list of floats.
|
||||
|
||||
Perplexity returns base64-encoded signed int8 values by default.
|
||||
If the value is already a list of numbers (e.g. from a mock or
|
||||
future float format), it is returned as-is.
|
||||
"""
|
||||
if isinstance(embedding_value, list):
|
||||
return embedding_value
|
||||
if isinstance(embedding_value, str):
|
||||
raw_bytes = base64.b64decode(embedding_value)
|
||||
count = len(raw_bytes)
|
||||
int8_values = struct.unpack(f"{count}b", raw_bytes)
|
||||
return [float(v) / 127.0 for v in int8_values]
|
||||
return embedding_value
|
||||
|
||||
def transform_embedding_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: EmbeddingResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: Optional[str] = None,
|
||||
request_data: dict = {},
|
||||
optional_params: dict = {},
|
||||
litellm_params: dict = {},
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
raw_response_json = raw_response.json()
|
||||
except Exception:
|
||||
raise PerplexityEmbeddingError(
|
||||
message=raw_response.text, status_code=raw_response.status_code
|
||||
)
|
||||
|
||||
model_response.model = raw_response_json.get("model", model)
|
||||
model_response.object = raw_response_json.get("object", "list")
|
||||
|
||||
raw_data = raw_response_json.get("data", [])
|
||||
decoded_data: List[Dict[str, Any]] = []
|
||||
for item in raw_data:
|
||||
decoded_item = dict(item)
|
||||
decoded_item["embedding"] = self._decode_base64_embedding(
|
||||
item.get("embedding")
|
||||
)
|
||||
decoded_data.append(decoded_item)
|
||||
model_response.data = decoded_data
|
||||
|
||||
usage_data = raw_response_json.get("usage", {})
|
||||
usage = Usage(
|
||||
prompt_tokens=usage_data.get("prompt_tokens", 0)
|
||||
or usage_data.get("total_tokens", 0),
|
||||
total_tokens=usage_data.get("total_tokens", 0),
|
||||
)
|
||||
model_response.usage = usage
|
||||
return model_response
|
||||
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: Union[dict, httpx.Headers],
|
||||
) -> BaseLLMException:
|
||||
return PerplexityEmbeddingError(
|
||||
message=error_message, status_code=status_code, headers=headers
|
||||
)
|
||||
@@ -5627,6 +5627,21 @@ def embedding( # noqa: PLR0915
|
||||
aembedding=aembedding,
|
||||
litellm_params={"ssl_verify": kwargs.get("ssl_verify", None)},
|
||||
)
|
||||
elif custom_llm_provider == "perplexity":
|
||||
response = base_llm_http_handler.embedding(
|
||||
model=model,
|
||||
input=input,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
logging_obj=logging,
|
||||
timeout=timeout,
|
||||
model_response=EmbeddingResponse(),
|
||||
optional_params=optional_params,
|
||||
client=client,
|
||||
aembedding=aembedding,
|
||||
litellm_params={},
|
||||
)
|
||||
else:
|
||||
raise LiteLLMUnknownProvider(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
|
||||
@@ -26952,6 +26952,26 @@
|
||||
"supports_reasoning": false,
|
||||
"supports_function_calling": true
|
||||
},
|
||||
"perplexity/pplx-embed-v1-0.6b": {
|
||||
"input_cost_per_token": 0.000000004,
|
||||
"litellm_provider": "perplexity",
|
||||
"max_input_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://docs.perplexity.ai/docs/embeddings/quickstart"
|
||||
},
|
||||
"perplexity/pplx-embed-v1-4b": {
|
||||
"input_cost_per_token": 0.00000003,
|
||||
"litellm_provider": "perplexity",
|
||||
"max_input_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 2560,
|
||||
"source": "https://docs.perplexity.ai/docs/embeddings/quickstart"
|
||||
},
|
||||
"publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "publicai",
|
||||
|
||||
@@ -8145,6 +8145,8 @@ class ProviderConfigManager:
|
||||
)
|
||||
|
||||
return SagemakerEmbeddingConfig.get_model_config(model)
|
||||
elif litellm.LlmProviders.PERPLEXITY == provider:
|
||||
return litellm.PerplexityEmbeddingConfig()
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -27187,6 +27187,26 @@
|
||||
"supports_reasoning": false,
|
||||
"supports_function_calling": true
|
||||
},
|
||||
"perplexity/pplx-embed-v1-0.6b": {
|
||||
"input_cost_per_token": 0.000000004,
|
||||
"litellm_provider": "perplexity",
|
||||
"max_input_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://docs.perplexity.ai/docs/embeddings/quickstart"
|
||||
},
|
||||
"perplexity/pplx-embed-v1-4b": {
|
||||
"input_cost_per_token": 0.00000003,
|
||||
"litellm_provider": "perplexity",
|
||||
"max_input_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 0.0,
|
||||
"output_vector_size": 2560,
|
||||
"source": "https://docs.perplexity.ai/docs/embeddings/quickstart"
|
||||
},
|
||||
"publicai/aisingapore/Qwen-SEA-LION-v4-32B-IT": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "publicai",
|
||||
|
||||
+320
@@ -0,0 +1,320 @@
|
||||
"""
|
||||
Unit tests for Perplexity embedding transformation logic.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import struct
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.perplexity.embedding.transformation import (
|
||||
PerplexityEmbeddingConfig,
|
||||
PerplexityEmbeddingError,
|
||||
)
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
class TestPerplexityEmbeddingConfig:
|
||||
def setup_method(self):
|
||||
self.config = PerplexityEmbeddingConfig()
|
||||
self.model = "pplx-embed-v1-0.6b"
|
||||
self.logging_obj = MagicMock()
|
||||
|
||||
def test_get_complete_url_default(self):
|
||||
"""Test default URL construction."""
|
||||
url = self.config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key="test-key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == "https://api.perplexity.ai/v1/embeddings"
|
||||
|
||||
def test_get_complete_url_custom_base(self):
|
||||
"""Test URL construction with custom api_base."""
|
||||
url = self.config.get_complete_url(
|
||||
api_base="https://custom.api.com",
|
||||
api_key="test-key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == "https://custom.api.com/v1/embeddings"
|
||||
|
||||
def test_get_complete_url_already_has_embeddings(self):
|
||||
"""Test URL construction when api_base already ends with /embeddings."""
|
||||
url = self.config.get_complete_url(
|
||||
api_base="https://custom.api.com/v1/embeddings",
|
||||
api_key="test-key",
|
||||
model=self.model,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert url == "https://custom.api.com/v1/embeddings"
|
||||
|
||||
def test_get_supported_openai_params(self):
|
||||
"""Test that supported params are correctly listed."""
|
||||
supported = self.config.get_supported_openai_params(self.model)
|
||||
assert "dimensions" in supported
|
||||
assert "encoding_format" in supported
|
||||
|
||||
def test_map_openai_params_dimensions(self):
|
||||
"""Test that dimensions parameter is correctly mapped."""
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params={"dimensions": 512},
|
||||
optional_params={},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
)
|
||||
assert result["dimensions"] == 512
|
||||
|
||||
def test_map_openai_params_encoding_format(self):
|
||||
"""Test that encoding_format parameter is correctly mapped."""
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params={"encoding_format": "base64_int8"},
|
||||
optional_params={},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
)
|
||||
assert result["encoding_format"] == "base64_int8"
|
||||
|
||||
def test_map_openai_params_unsupported_dropped(self):
|
||||
"""Test that unsupported parameters are not passed through."""
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params={"dimensions": 256, "user": "test-user"},
|
||||
optional_params={},
|
||||
model=self.model,
|
||||
drop_params=False,
|
||||
)
|
||||
assert result["dimensions"] == 256
|
||||
assert "user" not in result
|
||||
|
||||
def test_validate_environment_with_api_key(self):
|
||||
"""Test environment validation with explicit API key."""
|
||||
headers = self.config.validate_environment(
|
||||
headers={},
|
||||
model=self.model,
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="pplx-test-key",
|
||||
)
|
||||
assert headers["Authorization"] == "Bearer pplx-test-key"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
def test_transform_embedding_request_string_input(self):
|
||||
"""Test request transformation with string input."""
|
||||
result = self.config.transform_embedding_request(
|
||||
model=self.model,
|
||||
input="Hello world",
|
||||
optional_params={},
|
||||
headers={},
|
||||
)
|
||||
assert result["model"] == self.model
|
||||
assert result["input"] == "Hello world"
|
||||
|
||||
def test_transform_embedding_request_list_input(self):
|
||||
"""Test request transformation with list input."""
|
||||
input_data = ["Hello world", "Testing embeddings"]
|
||||
result = self.config.transform_embedding_request(
|
||||
model=self.model,
|
||||
input=input_data,
|
||||
optional_params={},
|
||||
headers={},
|
||||
)
|
||||
assert result["model"] == self.model
|
||||
assert result["input"] == input_data
|
||||
|
||||
def test_transform_embedding_request_with_params(self):
|
||||
"""Test request transformation with optional params."""
|
||||
result = self.config.transform_embedding_request(
|
||||
model=self.model,
|
||||
input=["Test"],
|
||||
optional_params={"dimensions": 256},
|
||||
headers={},
|
||||
)
|
||||
assert result["model"] == self.model
|
||||
assert result["input"] == ["Test"]
|
||||
assert result["dimensions"] == 256
|
||||
|
||||
def test_transform_embedding_response_float_passthrough(self):
|
||||
"""Test response transformation when embeddings are already float arrays."""
|
||||
mock_response_data = {
|
||||
"object": "list",
|
||||
"model": "pplx-embed-v1-0.6b",
|
||||
"data": [
|
||||
{
|
||||
"object": "embedding",
|
||||
"index": 0,
|
||||
"embedding": [0.1, 0.2, 0.3],
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 5,
|
||||
"total_tokens": 5,
|
||||
},
|
||||
}
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.json.return_value = mock_response_data
|
||||
mock_response.status_code = 200
|
||||
|
||||
model_response = EmbeddingResponse()
|
||||
result = self.config.transform_embedding_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
|
||||
assert result.model == "pplx-embed-v1-0.6b"
|
||||
assert result.object == "list"
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0]["embedding"] == [0.1, 0.2, 0.3]
|
||||
assert result.usage.prompt_tokens == 5
|
||||
assert result.usage.total_tokens == 5
|
||||
|
||||
def test_transform_embedding_response_base64_int8(self):
|
||||
"""Test decoding base64_int8 embeddings to float arrays (Perplexity default)."""
|
||||
int8_values = [127, -128, 0, 64, -64]
|
||||
b64_encoded = base64.b64encode(struct.pack(f"{len(int8_values)}b", *int8_values)).decode()
|
||||
|
||||
mock_response_data = {
|
||||
"object": "list",
|
||||
"model": "pplx-embed-v1-0.6b",
|
||||
"data": [
|
||||
{
|
||||
"object": "embedding",
|
||||
"index": 0,
|
||||
"embedding": b64_encoded,
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 3, "total_tokens": 3},
|
||||
}
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.json.return_value = mock_response_data
|
||||
mock_response.status_code = 200
|
||||
|
||||
model_response = EmbeddingResponse()
|
||||
result = self.config.transform_embedding_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
|
||||
embedding = result.data[0]["embedding"]
|
||||
assert isinstance(embedding, list)
|
||||
assert len(embedding) == 5
|
||||
assert all(isinstance(v, float) for v in embedding)
|
||||
assert abs(embedding[0] - 1.0) < 0.01
|
||||
assert abs(embedding[1] - (-128.0 / 127.0)) < 0.01
|
||||
assert embedding[2] == 0.0
|
||||
|
||||
def test_decode_base64_embedding_static(self):
|
||||
"""Test the static decode helper directly."""
|
||||
int8_values = [10, -10, 50, -50]
|
||||
b64_str = base64.b64encode(struct.pack("4b", *int8_values)).decode()
|
||||
result = PerplexityEmbeddingConfig._decode_base64_embedding(b64_str)
|
||||
assert len(result) == 4
|
||||
assert abs(result[0] - 10.0 / 127.0) < 1e-6
|
||||
assert abs(result[1] - (-10.0 / 127.0)) < 1e-6
|
||||
|
||||
def test_decode_base64_embedding_list_passthrough(self):
|
||||
"""Test that float lists pass through unchanged."""
|
||||
floats = [0.5, -0.3, 0.8]
|
||||
result = PerplexityEmbeddingConfig._decode_base64_embedding(floats)
|
||||
assert result == floats
|
||||
|
||||
def test_transform_embedding_response_error(self):
|
||||
"""Test that malformed response raises PerplexityEmbeddingError."""
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.json.side_effect = Exception("Invalid JSON")
|
||||
mock_response.text = "Server error"
|
||||
mock_response.status_code = 500
|
||||
|
||||
model_response = EmbeddingResponse()
|
||||
try:
|
||||
self.config.transform_embedding_response(
|
||||
model=self.model,
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
assert False, "Should have raised PerplexityEmbeddingError"
|
||||
except PerplexityEmbeddingError as e:
|
||||
assert e.status_code == 500
|
||||
assert "Server error" in e.message
|
||||
|
||||
def test_get_error_class(self):
|
||||
"""Test that get_error_class returns the correct error type."""
|
||||
error = self.config.get_error_class(
|
||||
error_message="Not found",
|
||||
status_code=404,
|
||||
headers={},
|
||||
)
|
||||
assert isinstance(error, PerplexityEmbeddingError)
|
||||
assert error.status_code == 404
|
||||
assert error.message == "Not found"
|
||||
|
||||
def test_transform_embedding_request_4b_model(self):
|
||||
"""Test request transformation with the 4b model."""
|
||||
model = "pplx-embed-v1-4b"
|
||||
result = self.config.transform_embedding_request(
|
||||
model=model,
|
||||
input=["Test text"],
|
||||
optional_params={"dimensions": 2560},
|
||||
headers={},
|
||||
)
|
||||
assert result["model"] == model
|
||||
assert result["dimensions"] == 2560
|
||||
|
||||
|
||||
class TestPerplexityEmbeddingProviderConfig:
|
||||
"""Test that Perplexity is correctly registered in ProviderConfigManager."""
|
||||
|
||||
def test_provider_config_returns_perplexity_embedding(self):
|
||||
import litellm
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
config = ProviderConfigManager.get_provider_embedding_config(
|
||||
model="pplx-embed-v1-0.6b",
|
||||
provider=litellm.LlmProviders.PERPLEXITY,
|
||||
)
|
||||
assert config is not None
|
||||
assert isinstance(config, PerplexityEmbeddingConfig)
|
||||
|
||||
def test_provider_config_returns_perplexity_embedding_4b(self):
|
||||
import litellm
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
config = ProviderConfigManager.get_provider_embedding_config(
|
||||
model="pplx-embed-v1-4b",
|
||||
provider=litellm.LlmProviders.PERPLEXITY,
|
||||
)
|
||||
assert config is not None
|
||||
assert isinstance(config, PerplexityEmbeddingConfig)
|
||||
|
||||
|
||||
class TestPerplexityEmbeddingModelInfo:
|
||||
"""Test that Perplexity embedding models are in model_prices_and_context_window."""
|
||||
|
||||
def test_model_info_available(self):
|
||||
import litellm
|
||||
|
||||
info = litellm.get_model_info("perplexity/pplx-embed-v1-0.6b")
|
||||
assert info is not None
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["max_input_tokens"] == 32768
|
||||
assert info["output_vector_size"] == 1024
|
||||
|
||||
def test_model_info_4b_available(self):
|
||||
import litellm
|
||||
|
||||
info = litellm.get_model_info("perplexity/pplx-embed-v1-4b")
|
||||
assert info is not None
|
||||
assert info["mode"] == "embedding"
|
||||
assert info["max_input_tokens"] == 32768
|
||||
assert info["output_vector_size"] == 2560
|
||||
Reference in New Issue
Block a user