perplexity-haystack 0.1.0__tar.gz → 0.1.1__tar.gz

perplexity_haystack-0.1.1/CHANGELOG.md

@@ -0,0 +1,9 @@
+# Changelog
+
+## [integrations/perplexity-v0.1.0] - 2026-05-13
+
+### 🌀 Miscellaneous
+
+- Add Perplexity integration package (#3262)
+
+<!-- generated by git-cliff -->

{perplexity_haystack-0.1.0 → perplexity_haystack-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: perplexity-haystack
-Version: 0.1.0
+Version: 0.1.1
 Summary: Haystack integration for Perplexity Web Search
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/perplexity#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues

{perplexity_haystack-0.1.0 → perplexity_haystack-0.1.1}/src/haystack_integrations/components/embedders/perplexity/document_embedder.py

@@ -3,11 +3,24 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import importlib.metadata
+from collections.abc import Iterable
 from typing import Any, ClassVar
 
-from haystack import component, default_from_dict, default_to_dict
+from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.components.embedders import OpenAIDocumentEmbedder
 from haystack.utils.auth import Secret
+from more_itertools import batched
+from openai import APIError
+from openai.types.create_embedding_response import CreateEmbeddingResponse
+from tqdm import tqdm
+from tqdm.asyncio import tqdm as async_tqdm
+
+from haystack_integrations.components.embedders.perplexity.embedding_encoding import (
+    _decode_embedding,
+    _validate_encoding_format,
+)
+
+logger = logging.getLogger(__name__)
 
 _INTEGRATION_SLUG = "haystack"
 _PACKAGE_NAME = "perplexity-haystack"
@@ -73,6 +86,7 @@ class PerplexityDocumentEmbedder(OpenAIDocumentEmbedder):
         progress_bar: bool = True,
         meta_fields_to_embed: list[str] | None = None,
         embedding_separator: str = "\n",
+        encoding_format: str = "base64_int8",
         timeout: float | None = None,
         max_retries: int | None = None,
         http_client_kwargs: dict[str, Any] | None = None,
@@ -99,6 +113,8 @@ class PerplexityDocumentEmbedder(OpenAIDocumentEmbedder):
             List of meta fields that should be embedded along with the Document text.
         :param embedding_separator:
             Separator used to concatenate the meta fields to the Document text.
+        :param encoding_format:
+            The Perplexity embedding encoding format. Supported values are `base64_int8` and `base64_binary`.
         :param timeout:
             Timeout for Perplexity client calls. If not set, it defaults to either the `OPENAI_TIMEOUT` environment
             variable, or 30 seconds.
@@ -109,6 +125,7 @@ class PerplexityDocumentEmbedder(OpenAIDocumentEmbedder):
             A dictionary of keyword arguments to configure a custom `httpx.Client`or `httpx.AsyncClient`.
             For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
         """
+        self.encoding_format = _validate_encoding_format(encoding_format)
         super(PerplexityDocumentEmbedder, self).__init__(  # noqa: UP008
             api_key=api_key,
             model=model,
@@ -129,6 +146,94 @@ class PerplexityDocumentEmbedder(OpenAIDocumentEmbedder):
         self.timeout = timeout
         self.max_retries = max_retries
 
+    def _decode_response_embeddings(self, response: CreateEmbeddingResponse) -> list[list[float]]:
+        return [_decode_embedding(str(el.embedding), self.encoding_format) for el in response.data]
+
+    def _embed_batch(
+        self, texts_to_embed: dict[str, str], batch_size: int
+    ) -> tuple[dict[str, list[float]], dict[str, Any]]:
+        """
+        Embed a list of texts in batches.
+        """
+
+        doc_ids_to_embeddings: dict[str, list[float]] = {}
+        meta: dict[str, Any] = {}
+        for batch in tqdm(
+            batched(texts_to_embed.items(), batch_size), disable=not self.progress_bar, desc="Calculating embeddings"
+        ):
+            args: dict[str, Any] = {
+                "model": self.model,
+                "input": [b[1] for b in batch],
+                "encoding_format": self.encoding_format,
+            }
+
+            try:
+                response = self.client.embeddings.create(**args)
+            except APIError as exc:
+                ids = ", ".join(b[0] for b in batch)
+                msg = "Failed embedding of documents {ids} caused by {exc}"
+                logger.exception(msg, ids=ids, exc=exc)
+                if self.raise_on_failure:
+                    raise exc
+                continue
+
+            embeddings = self._decode_response_embeddings(response)
+            doc_ids_to_embeddings.update(dict(zip((b[0] for b in batch), embeddings, strict=True)))
+
+            if "model" not in meta:
+                meta["model"] = response.model
+            if "usage" not in meta:
+                meta["usage"] = dict(response.usage)
+            else:
+                meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
+                meta["usage"]["total_tokens"] += response.usage.total_tokens
+
+        return doc_ids_to_embeddings, meta
+
+    async def _embed_batch_async(
+        self, texts_to_embed: dict[str, str], batch_size: int
+    ) -> tuple[dict[str, list[float]], dict[str, Any]]:
+        """
+        Embed a list of texts in batches asynchronously.
+        """
+
+        doc_ids_to_embeddings: dict[str, list[float]] = {}
+        meta: dict[str, Any] = {}
+
+        batches: Iterable[tuple[tuple[str, str], ...]] = list(batched(texts_to_embed.items(), batch_size))
+        if self.progress_bar:
+            batches = async_tqdm(batches, desc="Calculating embeddings")
+
+        for batch in batches:
+            args: dict[str, Any] = {
+                "model": self.model,
+                "input": [b[1] for b in batch],
+                "encoding_format": self.encoding_format,
+            }
+
+            try:
+                response = await self.async_client.embeddings.create(**args)
+            except APIError as exc:
+                ids = ", ".join(b[0] for b in batch)
+                msg = "Failed embedding of documents {ids} caused by {exc}"
+                logger.exception(msg, ids=ids, exc=exc)
+                if self.raise_on_failure:
+                    raise exc
+                continue
+
+            embeddings = self._decode_response_embeddings(response)
+            doc_ids_to_embeddings.update(dict(zip((b[0] for b in batch), embeddings, strict=True)))
+
+            if "model" not in meta:
+                meta["model"] = response.model
+            if "usage" not in meta:
+                meta["usage"] = dict(response.usage)
+            else:
+                meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens
+                meta["usage"]["total_tokens"] += response.usage.total_tokens
+
+        return doc_ids_to_embeddings, meta
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
@@ -147,6 +252,7 @@ class PerplexityDocumentEmbedder(OpenAIDocumentEmbedder):
             progress_bar=self.progress_bar,
             meta_fields_to_embed=self.meta_fields_to_embed,
             embedding_separator=self.embedding_separator,
+            encoding_format=self.encoding_format,
             timeout=self.timeout,
             max_retries=self.max_retries,
             http_client_kwargs=self.http_client_kwargs,

perplexity_haystack-0.1.1/src/haystack_integrations/components/embedders/perplexity/embedding_encoding.py

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import base64
+
+import numpy as np
+
+SUPPORTED_ENCODING_FORMATS = {"base64_int8", "base64_binary"}
+
+
+def _validate_encoding_format(encoding_format: str) -> str:
+    """
+    Validate Perplexity's embedding encoding format.
+    """
+    if encoding_format not in SUPPORTED_ENCODING_FORMATS:
+        supported_formats = "', '".join(sorted(SUPPORTED_ENCODING_FORMATS))
+        msg = f"Unsupported encoding_format='{encoding_format}'. Use '{supported_formats}'."
+        raise ValueError(msg)
+    return encoding_format
+
+
+def _decode_embedding(embedding: str, encoding_format: str) -> list[float]:
+    """
+    Decode a Perplexity base64 embedding into Haystack's list[float] representation.
+    """
+    raw_embedding = base64.b64decode(embedding)
+    if encoding_format == "base64_int8":
+        return np.frombuffer(raw_embedding, dtype=np.int8).astype(np.float32).tolist()
+    if encoding_format == "base64_binary":
+        return np.unpackbits(np.frombuffer(raw_embedding, dtype=np.uint8)).astype(np.float32).tolist()
+
+    msg = f"Unsupported encoding_format='{encoding_format}'."
+    raise ValueError(msg)

{perplexity_haystack-0.1.0 → perplexity_haystack-0.1.1}/src/haystack_integrations/components/embedders/perplexity/text_embedder.py

@@ -8,6 +8,12 @@ from typing import Any, ClassVar
 from haystack import component, default_from_dict, default_to_dict
 from haystack.components.embedders import OpenAITextEmbedder
 from haystack.utils.auth import Secret
+from openai.types.create_embedding_response import CreateEmbeddingResponse
+
+from haystack_integrations.components.embedders.perplexity.embedding_encoding import (
+    _decode_embedding,
+    _validate_encoding_format,
+)
 
 _INTEGRATION_SLUG = "haystack"
 _PACKAGE_NAME = "perplexity-haystack"
@@ -64,6 +70,7 @@ class PerplexityTextEmbedder(OpenAITextEmbedder):
         api_base_url: str | None = "https://api.perplexity.ai/v1",
         prefix: str = "",
         suffix: str = "",
+        encoding_format: str = "base64_int8",
         timeout: float | None = None,
         max_retries: int | None = None,
         http_client_kwargs: dict[str, Any] | None = None,
@@ -81,6 +88,8 @@ class PerplexityTextEmbedder(OpenAITextEmbedder):
             A string to add to the beginning of each text.
         :param suffix:
             A string to add to the end of each text.
+        :param encoding_format:
+            The Perplexity embedding encoding format. Supported values are `base64_int8` and `base64_binary`.
         :param timeout:
             Timeout for Perplexity client calls. If not set, it defaults to either the `OPENAI_TIMEOUT` environment
             variable, or 30 seconds.
@@ -92,6 +101,7 @@ class PerplexityTextEmbedder(OpenAITextEmbedder):
             For more information, see the [HTTPX documentation](https://www.python-httpx.org/api/#client).
         """
 
+        self.encoding_format = _validate_encoding_format(encoding_format)
         super(PerplexityTextEmbedder, self).__init__(  # noqa: UP008
             api_key=api_key,
             model=model,
@@ -108,6 +118,18 @@ class PerplexityTextEmbedder(OpenAITextEmbedder):
         self.timeout = timeout
         self.max_retries = max_retries
 
+    def _prepare_input(self, text: str) -> dict[str, Any]:
+        kwargs = OpenAITextEmbedder._prepare_input(self, text=text)
+        kwargs["input"] = [kwargs["input"]]
+        kwargs["encoding_format"] = self.encoding_format
+        return kwargs
+
+    def _prepare_output(self, result: CreateEmbeddingResponse) -> dict[str, Any]:
+        return {
+            "embedding": _decode_embedding(str(result.data[0].embedding), self.encoding_format),
+            "meta": {"model": result.model, "usage": dict(result.usage)},
+        }
+
     def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
@@ -122,6 +144,7 @@ class PerplexityTextEmbedder(OpenAITextEmbedder):
             api_base_url=self.api_base_url,
             prefix=self.prefix,
             suffix=self.suffix,
+            encoding_format=self.encoding_format,
             timeout=self.timeout,
             max_retries=self.max_retries,
             http_client_kwargs=self.http_client_kwargs,

perplexity_haystack-0.1.1/tests/test_embedding_encoding.py

@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from haystack_integrations.components.embedders.perplexity.embedding_encoding import (
+    _decode_embedding,
+    _validate_encoding_format,
+)
+
+
+def test_validate_encoding_format_rejects_unsupported_format():
+    with pytest.raises(ValueError) as exc_info:
+        _validate_encoding_format("base64_float16")
+
+    assert str(exc_info.value) == "Unsupported encoding_format='base64_float16'. Use 'base64_binary', 'base64_int8'."
+
+
+def test_decode_embedding_rejects_unsupported_format():
+    with pytest.raises(ValueError) as exc_info:
+        _decode_embedding("", "base64_float16")
+
+    assert str(exc_info.value) == "Unsupported encoding_format='base64_float16'."

{perplexity_haystack-0.1.0 → perplexity_haystack-0.1.1}/tests/test_perplexity_document_embedder.py

@@ -2,13 +2,17 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import base64
 import json
 import os
+from types import SimpleNamespace
 
 import httpx
+import numpy as np
 import pytest
 from haystack import Document
 from haystack.utils import Secret
+from openai import APIError
 
 from haystack_integrations.components.embedders.perplexity import (
     document_embedder as document_embedder_module,
@@ -17,17 +21,24 @@ from haystack_integrations.components.embedders.perplexity.document_embedder imp
     PerplexityDocumentEmbedder,
 )
 
+INT8_EMBEDDING = [-2.0, -1.0, 0.0, 1.0, 2.0]
+INT8_EMBEDDING_PAYLOAD = base64.b64encode(np.array(INT8_EMBEDDING, dtype=np.int8).tobytes()).decode()
+BINARY_EMBEDDING = [1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]
+BINARY_EMBEDDING_PAYLOAD = base64.b64encode(np.packbits(np.array(BINARY_EMBEDDING, dtype=np.uint8)).tobytes()).decode()
 
-def _make_transport(captured: list[httpx.Request]) -> httpx.MockTransport:
+
+def _make_transport(captured: list[httpx.Request], embeddings: list[str] | None = None) -> httpx.MockTransport:
     def handler(request: httpx.Request) -> httpx.Response:
         captured.append(request)
+        inputs = json.loads(request.content)["input"]
+        response_embeddings = embeddings or [INT8_EMBEDDING_PAYLOAD] * len(inputs)
         return httpx.Response(
             200,
             json={
                 "object": "list",
                 "data": [
-                    {"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0},
-                    {"object": "embedding", "embedding": [0.4, 0.5, 0.6], "index": 1},
+                    {"object": "embedding", "embedding": embedding, "index": index}
+                    for index, embedding in enumerate(response_embeddings)
                 ],
                 "model": "pplx-embed-v1-0.6b",
                 "usage": {"prompt_tokens": 6, "total_tokens": 6},
@@ -38,6 +49,14 @@ def _make_transport(captured: list[httpx.Request]) -> httpx.MockTransport:
     return httpx.MockTransport(handler)
 
 
+def _make_api_error() -> APIError:
+    return APIError(
+        "Failed embedding",
+        request=httpx.Request("POST", "https://api.perplexity.ai/v1/embeddings"),
+        body=None,
+    )
+
+
 class TestPerplexityDocumentEmbedder:
     def test_attribution_header_falls_back_when_package_is_not_installed(self, monkeypatch):
         def raise_package_not_found(_package_name: str) -> str:
@@ -77,6 +96,7 @@ class TestPerplexityDocumentEmbedder:
         assert embedder.progress_bar is True
         assert embedder.meta_fields_to_embed == []
         assert embedder.embedding_separator == "\n"
+        assert embedder.encoding_format == "base64_int8"
 
     def test_init_with_parameters(self):
         embedder = PerplexityDocumentEmbedder(
@@ -89,6 +109,7 @@ class TestPerplexityDocumentEmbedder:
             progress_bar=False,
             meta_fields_to_embed=["test_field"],
             embedding_separator="-",
+            encoding_format="base64_binary",
         )
 
         assert embedder.api_key == Secret.from_token("test-api-key")
@@ -100,6 +121,14 @@ class TestPerplexityDocumentEmbedder:
         assert embedder.progress_bar is False
         assert embedder.meta_fields_to_embed == ["test_field"]
         assert embedder.embedding_separator == "-"
+        assert embedder.encoding_format == "base64_binary"
+
+    def test_init_rejects_float_encoding_format(self):
+        with pytest.raises(ValueError):
+            PerplexityDocumentEmbedder(
+                api_key=Secret.from_token("test-api-key"),
+                encoding_format="float",
+            )
 
     def test_to_dict(self, monkeypatch):
         monkeypatch.setenv("PERPLEXITY_API_KEY", "test-api-key")
@@ -122,6 +151,7 @@ class TestPerplexityDocumentEmbedder:
                 "progress_bar": True,
                 "meta_fields_to_embed": [],
                 "embedding_separator": "\n",
+                "encoding_format": "base64_int8",
                 "timeout": None,
                 "max_retries": None,
                 "http_client_kwargs": None,
@@ -140,6 +170,7 @@ class TestPerplexityDocumentEmbedder:
             progress_bar=False,
             meta_fields_to_embed=["test_field"],
             embedding_separator="-",
+            encoding_format="base64_binary",
             timeout=10.0,
             max_retries=2,
             http_client_kwargs={"proxy": "http://localhost:8080"},
@@ -161,6 +192,7 @@ class TestPerplexityDocumentEmbedder:
                 "progress_bar": False,
                 "meta_fields_to_embed": ["test_field"],
                 "embedding_separator": "-",
+                "encoding_format": "base64_binary",
                 "timeout": 10.0,
                 "max_retries": 2,
                 "http_client_kwargs": {"proxy": "http://localhost:8080"},
@@ -183,6 +215,7 @@ class TestPerplexityDocumentEmbedder:
                 "progress_bar": True,
                 "meta_fields_to_embed": [],
                 "embedding_separator": "\n",
+                "encoding_format": "base64_binary",
                 "timeout": None,
                 "max_retries": None,
                 "http_client_kwargs": None,
@@ -200,6 +233,7 @@ class TestPerplexityDocumentEmbedder:
         assert component.progress_bar is True
         assert component.meta_fields_to_embed == []
         assert component.embedding_separator == "\n"
+        assert component.encoding_format == "base64_binary"
         assert component.timeout is None
         assert component.max_retries is None
         assert component.http_client_kwargs is None
@@ -222,8 +256,8 @@ class TestPerplexityDocumentEmbedder:
         result = embedder.run(docs)
 
         docs_with_embeddings = result["documents"]
-        assert docs_with_embeddings[0].embedding == [0.1, 0.2, 0.3]
-        assert docs_with_embeddings[1].embedding == [0.4, 0.5, 0.6]
+        assert docs_with_embeddings[0].embedding == INT8_EMBEDDING
+        assert docs_with_embeddings[1].embedding == INT8_EMBEDDING
         assert len(captured) == 1
         request = captured[0]
         assert request.headers["Authorization"] == "Bearer test-api-key"
@@ -234,7 +268,172 @@ class TestPerplexityDocumentEmbedder:
             "I love cheese",
             "A transformer is a deep learning architecture",
         ]
-        assert body["encoding_format"] == "float"
+        assert body["encoding_format"] == "base64_int8"
+
+    @pytest.mark.parametrize(
+        ("encoding_format", "payload", "expected_embedding"),
+        [
+            ("base64_int8", INT8_EMBEDDING_PAYLOAD, INT8_EMBEDDING),
+            ("base64_binary", BINARY_EMBEDDING_PAYLOAD, BINARY_EMBEDDING),
+        ],
+    )
+    def test_run_decodes_base64_embeddings(self, encoding_format, payload, expected_embedding):
+        captured: list[httpx.Request] = []
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            encoding_format=encoding_format,
+            progress_bar=False,
+            http_client_kwargs={"transport": _make_transport(captured, embeddings=[payload])},
+        )
+
+        result = embedder.run([Document(content="I love cheese")])
+
+        embedding = result["documents"][0].embedding
+        assert embedding == expected_embedding
+        assert len(embedding) == len(expected_embedding)
+        assert all(isinstance(value, float) for value in embedding)
+        body = json.loads(captured[0].content)
+        assert body["encoding_format"] == encoding_format
+
+    def test_run_accumulates_usage_across_batches(self):
+        captured: list[httpx.Request] = []
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            batch_size=1,
+            progress_bar=False,
+            http_client_kwargs={"transport": _make_transport(captured)},
+        )
+
+        result = embedder.run(
+            [
+                Document(content="I love cheese"),
+                Document(content="A transformer is a deep learning architecture"),
+            ]
+        )
+
+        assert [document.embedding for document in result["documents"]] == [INT8_EMBEDDING, INT8_EMBEDDING]
+        assert result["meta"]["usage"] == {"prompt_tokens": 12, "total_tokens": 12}
+        assert len(captured) == 2
+
+    def test_embed_batch_skips_api_error_when_raise_on_failure_is_false(self):
+        api_error = _make_api_error()
+
+        def create(**_kwargs):
+            raise api_error
+
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            progress_bar=False,
+        )
+        embedder.raise_on_failure = False
+        embedder.client = SimpleNamespace(embeddings=SimpleNamespace(create=create))
+
+        embeddings, meta = embedder._embed_batch({"doc-1": "I love cheese"}, batch_size=1)
+
+        assert embeddings == {}
+        assert meta == {}
+
+    def test_embed_batch_raises_api_error_when_raise_on_failure_is_true(self):
+        api_error = _make_api_error()
+
+        def create(**_kwargs):
+            raise api_error
+
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            progress_bar=False,
+        )
+        embedder.raise_on_failure = True
+        embedder.client = SimpleNamespace(embeddings=SimpleNamespace(create=create))
+
+        with pytest.raises(APIError) as exc_info:
+            embedder._embed_batch({"doc-1": "I love cheese"}, batch_size=1)
+
+        assert exc_info.value is api_error
+
+    @pytest.mark.asyncio
+    async def test_run_async_decodes_base64_embeddings(self):
+        captured: list[httpx.Request] = []
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            progress_bar=False,
+            http_client_kwargs={"transport": _make_transport(captured, embeddings=[BINARY_EMBEDDING_PAYLOAD])},
+            encoding_format="base64_binary",
+        )
+
+        result = await embedder.run_async([Document(content="I love cheese")])
+
+        assert result["documents"][0].embedding == BINARY_EMBEDDING
+        body = json.loads(captured[0].content)
+        assert body["encoding_format"] == "base64_binary"
+
+    @pytest.mark.asyncio
+    async def test_run_async_uses_progress_bar_and_accumulates_usage_across_batches(self, monkeypatch):
+        captured: list[httpx.Request] = []
+        batches_seen = []
+
+        def fake_async_tqdm(batches, desc):
+            assert desc == "Calculating embeddings"
+            batches_seen.extend(batches)
+            return batches
+
+        monkeypatch.setattr(document_embedder_module, "async_tqdm", fake_async_tqdm)
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            batch_size=1,
+            progress_bar=True,
+            http_client_kwargs={"transport": _make_transport(captured)},
+        )
+
+        result = await embedder.run_async(
+            [
+                Document(content="I love cheese"),
+                Document(content="A transformer is a deep learning architecture"),
+            ]
+        )
+
+        assert [document.embedding for document in result["documents"]] == [INT8_EMBEDDING, INT8_EMBEDDING]
+        assert result["meta"]["usage"] == {"prompt_tokens": 12, "total_tokens": 12}
+        assert len(batches_seen) == 2
+        assert len(captured) == 2
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_async_skips_api_error_when_raise_on_failure_is_false(self):
+        api_error = _make_api_error()
+
+        async def create(**_kwargs):
+            raise api_error
+
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            progress_bar=False,
+        )
+        embedder.raise_on_failure = False
+        embedder.async_client = SimpleNamespace(embeddings=SimpleNamespace(create=create))
+
+        embeddings, meta = await embedder._embed_batch_async({"doc-1": "I love cheese"}, batch_size=1)
+
+        assert embeddings == {}
+        assert meta == {}
+
+    @pytest.mark.asyncio
+    async def test_embed_batch_async_raises_api_error_when_raise_on_failure_is_true(self):
+        api_error = _make_api_error()
+
+        async def create(**_kwargs):
+            raise api_error
+
+        embedder = PerplexityDocumentEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            progress_bar=False,
+        )
+        embedder.raise_on_failure = True
+        embedder.async_client = SimpleNamespace(embeddings=SimpleNamespace(create=create))
+
+        with pytest.raises(APIError) as exc_info:
+            await embedder._embed_batch_async({"doc-1": "I love cheese"}, batch_size=1)
+
+        assert exc_info.value is api_error
 
     def test_run_wrong_input_format(self):
         embedder = PerplexityDocumentEmbedder(api_key=Secret.from_token("test-api-key"))
@@ -269,7 +468,7 @@ class TestPerplexityDocumentEmbedderInference:
         assert len(embedded_docs) == len(docs)
         for doc in embedded_docs:
             assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) > 0
+            assert len(doc.embedding) == 1024
             assert all(isinstance(x, float) for x in doc.embedding)
 
     @pytest.mark.asyncio
@@ -285,4 +484,5 @@ class TestPerplexityDocumentEmbedderInference:
         assert len(embedded_docs) == len(docs)
         for doc in embedded_docs:
             assert isinstance(doc.embedding, list)
-            assert len(doc.embedding) > 0
+            assert len(doc.embedding) == 1024
+            assert all(isinstance(x, float) for x in doc.embedding)

{perplexity_haystack-0.1.0 → perplexity_haystack-0.1.1}/tests/test_perplexity_text_embedder.py

@@ -2,10 +2,12 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import base64
 import json
 import os
 
 import httpx
+import numpy as np
 import pytest
 from haystack.utils import Secret
 
@@ -16,15 +18,20 @@ from haystack_integrations.components.embedders.perplexity.text_embedder import
     PerplexityTextEmbedder,
 )
 
+INT8_EMBEDDING = [-2.0, -1.0, 0.0, 1.0, 2.0]
+INT8_EMBEDDING_PAYLOAD = base64.b64encode(np.array(INT8_EMBEDDING, dtype=np.int8).tobytes()).decode()
+BINARY_EMBEDDING = [1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 1.0]
+BINARY_EMBEDDING_PAYLOAD = base64.b64encode(np.packbits(np.array(BINARY_EMBEDDING, dtype=np.uint8)).tobytes()).decode()
 
-def _make_transport(captured: list[httpx.Request]) -> httpx.MockTransport:
+
+def _make_transport(captured: list[httpx.Request], embedding: str = INT8_EMBEDDING_PAYLOAD) -> httpx.MockTransport:
     def handler(request: httpx.Request) -> httpx.Response:
         captured.append(request)
         return httpx.Response(
             200,
             json={
                 "object": "list",
-                "data": [{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}],
+                "data": [{"object": "embedding", "embedding": embedding, "index": 0}],
                 "model": "pplx-embed-v1-0.6b",
                 "usage": {"prompt_tokens": 3, "total_tokens": 3},
             },
@@ -67,6 +74,7 @@ class TestPerplexityTextEmbedder:
         assert embedder.model == "pplx-embed-v1-0.6b"
         assert embedder.prefix == ""
         assert embedder.suffix == ""
+        assert embedder.encoding_format == "base64_int8"
 
     def test_init_with_parameters(self):
         embedder = PerplexityTextEmbedder(
@@ -74,6 +82,7 @@ class TestPerplexityTextEmbedder:
             model="pplx-embed-v1-4b",
             prefix="START",
             suffix="END",
+            encoding_format="base64_binary",
         )
 
         assert embedder.api_key == Secret.from_token("test-api-key")
@@ -81,6 +90,14 @@ class TestPerplexityTextEmbedder:
         assert embedder.model == "pplx-embed-v1-4b"
         assert embedder.prefix == "START"
         assert embedder.suffix == "END"
+        assert embedder.encoding_format == "base64_binary"
+
+    def test_init_rejects_float_encoding_format(self):
+        with pytest.raises(ValueError):
+            PerplexityTextEmbedder(
+                api_key=Secret.from_token("test-api-key"),
+                encoding_format="float",
+            )
 
     def test_to_dict(self, monkeypatch):
         monkeypatch.setenv("PERPLEXITY_API_KEY", "test-api-key")
@@ -97,6 +114,7 @@ class TestPerplexityTextEmbedder:
                 "api_base_url": "https://api.perplexity.ai/v1",
                 "prefix": "",
                 "suffix": "",
+                "encoding_format": "base64_int8",
                 "timeout": None,
                 "max_retries": None,
                 "http_client_kwargs": None,
@@ -111,6 +129,7 @@ class TestPerplexityTextEmbedder:
             api_base_url="https://custom-api-base-url.com",
             prefix="START",
             suffix="END",
+            encoding_format="base64_binary",
             timeout=10.0,
             max_retries=2,
             http_client_kwargs={"proxy": "http://localhost:8080"},
@@ -126,6 +145,7 @@ class TestPerplexityTextEmbedder:
                 "api_base_url": "https://custom-api-base-url.com",
                 "prefix": "START",
                 "suffix": "END",
+                "encoding_format": "base64_binary",
                 "timeout": 10.0,
                 "max_retries": 2,
                 "http_client_kwargs": {"proxy": "http://localhost:8080"},
@@ -142,6 +162,7 @@ class TestPerplexityTextEmbedder:
                 "api_base_url": "https://api.perplexity.ai/v1",
                 "prefix": "",
                 "suffix": "",
+                "encoding_format": "base64_binary",
                 "timeout": None,
                 "max_retries": None,
                 "http_client_kwargs": None,
@@ -155,6 +176,7 @@ class TestPerplexityTextEmbedder:
         assert component.model == "pplx-embed-v1-0.6b"
         assert component.prefix == ""
         assert component.suffix == ""
+        assert component.encoding_format == "base64_binary"
         assert component.timeout is None
         assert component.max_retries is None
         assert component.http_client_kwargs is None
@@ -168,15 +190,54 @@ class TestPerplexityTextEmbedder:
 
         result = embedder.run(text="The food was delicious")
 
-        assert result["embedding"] == [0.1, 0.2, 0.3]
+        assert result["embedding"] == INT8_EMBEDDING
         assert len(captured) == 1
         request = captured[0]
         assert request.headers["Authorization"] == "Bearer test-api-key"
         assert request.headers["X-Pplx-Integration"].startswith("haystack/")
         body = json.loads(request.content)
         assert body["model"] == "pplx-embed-v1-0.6b"
-        assert body["input"] == "The food was delicious"
-        assert body["encoding_format"] == "float"
+        assert body["input"] == ["The food was delicious"]
+        assert body["encoding_format"] == "base64_int8"
+
+    @pytest.mark.parametrize(
+        ("encoding_format", "payload", "expected_embedding"),
+        [
+            ("base64_int8", INT8_EMBEDDING_PAYLOAD, INT8_EMBEDDING),
+            ("base64_binary", BINARY_EMBEDDING_PAYLOAD, BINARY_EMBEDDING),
+        ],
+    )
+    def test_run_decodes_base64_embeddings(self, encoding_format, payload, expected_embedding):
+        captured: list[httpx.Request] = []
+        embedder = PerplexityTextEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            encoding_format=encoding_format,
+            http_client_kwargs={"transport": _make_transport(captured, embedding=payload)},
+        )
+
+        result = embedder.run(text="The food was delicious")
+
+        embedding = result["embedding"]
+        assert embedding == expected_embedding
+        assert len(embedding) == len(expected_embedding)
+        assert all(isinstance(value, float) for value in embedding)
+        body = json.loads(captured[0].content)
+        assert body["encoding_format"] == encoding_format
+
+    @pytest.mark.asyncio
+    async def test_run_async_decodes_base64_embeddings(self):
+        captured: list[httpx.Request] = []
+        embedder = PerplexityTextEmbedder(
+            api_key=Secret.from_token("test-api-key"),
+            encoding_format="base64_binary",
+            http_client_kwargs={"transport": _make_transport(captured, embedding=BINARY_EMBEDDING_PAYLOAD)},
+        )
+
+        result = await embedder.run_async(text="The food was delicious")
+
+        assert result["embedding"] == BINARY_EMBEDDING
+        body = json.loads(captured[0].content)
+        assert body["encoding_format"] == "base64_binary"
 
     def test_run_wrong_input_format(self):
         embedder = PerplexityTextEmbedder(api_key=Secret.from_token("test-api-key"))
@@ -201,7 +262,7 @@ class TestPerplexityTextEmbedderInference:
         result = embedder.run(text=text)
 
         assert isinstance(result["embedding"], list)
-        assert len(result["embedding"]) > 0
+        assert len(result["embedding"]) == 1024
         assert all(isinstance(x, float) for x in result["embedding"])
 
     @pytest.mark.asyncio
@@ -211,4 +272,5 @@ class TestPerplexityTextEmbedderInference:
         result = await embedder.run_async(text=text)
 
         assert isinstance(result["embedding"], list)
-        assert len(result["embedding"]) > 0
+        assert len(result["embedding"]) == 1024
+        assert all(isinstance(x, float) for x in result["embedding"])