From 1d9746c8fd1c73c2e7648c6da84f414c61e9ce8d Mon Sep 17 00:00:00 2001 From: Alejandro Do Nascimento Mora Date: Mon, 30 Jun 2025 12:22:20 +0200 Subject: [PATCH] fix: pass voyage api key to client --- projects/pgai/pgai/vectorizer/embedders/litellm.py | 2 +- projects/pgai/pgai/vectorizer/embedders/voyageai.py | 8 +++++--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/projects/pgai/pgai/vectorizer/embedders/litellm.py b/projects/pgai/pgai/vectorizer/embedders/litellm.py index c57f09728..01ecc3ef1 100644 --- a/projects/pgai/pgai/vectorizer/embedders/litellm.py +++ b/projects/pgai/pgai/vectorizer/embedders/litellm.py @@ -152,7 +152,7 @@ def token_counter(text: str) -> int: return token_counter case "voyage": - return voyage_token_counter(model) + return voyage_token_counter(model, self._api_key) case _: return None diff --git a/projects/pgai/pgai/vectorizer/embedders/voyageai.py b/projects/pgai/pgai/vectorizer/embedders/voyageai.py index 3e10ee67f..f851cb104 100644 --- a/projects/pgai/pgai/vectorizer/embedders/voyageai.py +++ b/projects/pgai/pgai/vectorizer/embedders/voyageai.py @@ -30,11 +30,13 @@ def voyage_max_tokens_per_batch(model: str) -> int: return 120_000 # NOTE: This is conservative, but there probably won't be new Voyage models, so... -def voyage_token_counter(model: str) -> Callable[[str], int] | None: +def voyage_token_counter( + model: str, api_key: str | None = None +) -> Callable[[str], int] | None: # Note: deferred import to avoid import overhead import voyageai - client: voyageai.Client = voyageai.Client() + client: voyageai.Client = voyageai.Client(api_key=api_key) try: tokenizer: Tokenizer = client.tokenizer(model) return lambda text: len(tokenizer.encode(text).tokens) @@ -91,7 +93,7 @@ def _max_tokens_per_batch(self) -> int | None: return voyage_max_tokens_per_batch(self.model) def _token_counter(self) -> Callable[[str], int] | None: - return voyage_token_counter(self.model) + return voyage_token_counter(self.model, self._api_key) @override async def call_embed_api(self, documents: list[str]) -> EmbeddingResponse: