From bfceb7fc3f71f2c1ad2820668067afc7fdbf64ef Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 2 Mar 2026 17:37:50 -0800 Subject: [PATCH] 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 * 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 * 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 --------- Co-authored-by: Cursor Agent Co-authored-by: Ishaan Jaff --- .../docs/providers/perplexity_embedding.md | 134 ++++++++ docs/my-website/sidebars.js | 9 +- litellm/__init__.py | 1 + litellm/_lazy_imports_registry.py | 5 + litellm/llms/perplexity/embedding/__init__.py | 0 .../perplexity/embedding/transformation.py | 189 +++++++++++ litellm/main.py | 15 + ...odel_prices_and_context_window_backup.json | 20 ++ litellm/utils.py | 2 + model_prices_and_context_window.json | 20 ++ .../test_litellm/llms/perplexity/__init__.py | 0 .../llms/perplexity/embedding/__init__.py | 0 ...est_perplexity_embedding_transformation.py | 320 ++++++++++++++++++ 13 files changed, 714 insertions(+), 1 deletion(-) create mode 100644 docs/my-website/docs/providers/perplexity_embedding.md create mode 100644 litellm/llms/perplexity/embedding/__init__.py create mode 100644 litellm/llms/perplexity/embedding/transformation.py create mode 100644 tests/test_litellm/llms/perplexity/__init__.py create mode 100644 tests/test_litellm/llms/perplexity/embedding/__init__.py create mode 100644 tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py diff --git a/docs/my-website/docs/providers/perplexity_embedding.md b/docs/my-website/docs/providers/perplexity_embedding.md new file mode 100644 index 0000000000..92981b2632 --- /dev/null +++ b/docs/my-website/docs/providers/perplexity_embedding.md @@ -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 + + + + +```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) +``` + + + + +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"] + }' +``` + + + + +## 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 + + + + +```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'])}") +``` + + + + +```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 + }' +``` + + + + +## Supported Models + +All models listed on the [Perplexity Embeddings docs](https://docs.perplexity.ai/docs/embeddings/quickstart) are supported. Use `model=perplexity/`. + +| 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 diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index a8580b183a..004114c8e0 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -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", diff --git a/litellm/__init__.py b/litellm/__init__.py index 59b8e2da2a..f00b816be5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 554827b7bc..6ff997b453 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -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", diff --git a/litellm/llms/perplexity/embedding/__init__.py b/litellm/llms/perplexity/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py new file mode 100644 index 0000000000..24881ccebf --- /dev/null +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -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 + ) diff --git a/litellm/main.py b/litellm/main.py index 378de17396..c3ac4c24ae 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cbd64a178b..48a93830d4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -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", diff --git a/litellm/utils.py b/litellm/utils.py index 400ea40d65..cbe6aa8e79 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8145,6 +8145,8 @@ class ProviderConfigManager: ) return SagemakerEmbeddingConfig.get_model_config(model) + elif litellm.LlmProviders.PERPLEXITY == provider: + return litellm.PerplexityEmbeddingConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4c0694db37..f2aa6287a5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -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", diff --git a/tests/test_litellm/llms/perplexity/__init__.py b/tests/test_litellm/llms/perplexity/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/perplexity/embedding/__init__.py b/tests/test_litellm/llms/perplexity/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py new file mode 100644 index 0000000000..c2dae49ece --- /dev/null +++ b/tests/test_litellm/llms/perplexity/embedding/test_perplexity_embedding_transformation.py @@ -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