From 112f82e2b2e3e106126a96e2fe3ac40197769ed2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 5 Dec 2025 06:55:09 +0000 Subject: [PATCH 1/3] Initial plan From 8727b8030d89db1208d51a39999e92ae386b7ae6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 5 Dec 2025 07:01:03 +0000 Subject: [PATCH 2/3] feat: rebase patch-1 changes on main - add gemini support and env var config Co-authored-by: NickCrews <10820686+NickCrews@users.noreply.github.com> --- projects/pgai/pgai/vectorizer/embedders/litellm.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/projects/pgai/pgai/vectorizer/embedders/litellm.py b/projects/pgai/pgai/vectorizer/embedders/litellm.py index 01ecc3ef1..da8f650aa 100644 --- a/projects/pgai/pgai/vectorizer/embedders/litellm.py +++ b/projects/pgai/pgai/vectorizer/embedders/litellm.py @@ -1,4 +1,5 @@ from collections.abc import AsyncGenerator, Callable +import os from typing import Any, Literal from pydantic import BaseModel @@ -61,6 +62,9 @@ def _max_chunks_per_batch(self) -> int: # Note: deferred import to avoid import overhead import litellm + if result := os.getenv("PGAI_MAX_CHUNKS_PER_BATCH"): + return int(result) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(self.model) # type: ignore match custom_llm_provider: case "cohere": @@ -71,6 +75,8 @@ def _max_chunks_per_batch(self) -> int: return 2048 # https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/embeddings?tabs=console#verify-inputs-dont-exceed-the-maximum-length case "bedrock": return 96 # NOTE: currently (Jan 2025) Bedrock only supports embeddings with Cohere or Titan models. The Titan API only processes one input per request, which LiteLLM already handles under the hood. We assume that the Cohere API has the same input limits as above. + case "gemini": + return 250 # https://docs.cloud.google.com/vertex-ai/docs/quotas#text-embedding-limits case "huggingface": return 2048 # NOTE: There is not documented limit. In testing we got a response for a request with 10k (short) inputs. case "mistral": @@ -80,10 +86,11 @@ def _max_chunks_per_batch(self) -> int: case "voyage": return 128 # see https://docs.voyageai.com/reference/embeddings-api case _: + fallback = 5 logger.warn( - f"unknown provider '{custom_llm_provider}', falling back to conservative max chunks per batch" + f"unknown provider '{custom_llm_provider}', falling back to {fallback} max chunks per batch" ) - return 5 + return fallback @override def _max_tokens_per_batch(self) -> int | None: From 4faf536dba27b7e33da7fefdbca1ce2eef0e5296 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 5 Dec 2025 07:02:22 +0000 Subject: [PATCH 3/3] fix: add validation for PGAI_MAX_CHUNKS_PER_BATCH env var Co-authored-by: NickCrews <10820686+NickCrews@users.noreply.github.com> --- projects/pgai/pgai/vectorizer/embedders/litellm.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/projects/pgai/pgai/vectorizer/embedders/litellm.py b/projects/pgai/pgai/vectorizer/embedders/litellm.py index da8f650aa..fadbe20bc 100644 --- a/projects/pgai/pgai/vectorizer/embedders/litellm.py +++ b/projects/pgai/pgai/vectorizer/embedders/litellm.py @@ -63,7 +63,18 @@ def _max_chunks_per_batch(self) -> int: import litellm if result := os.getenv("PGAI_MAX_CHUNKS_PER_BATCH"): - return int(result) + try: + max_chunks = int(result) + if max_chunks <= 0: + logger.warn( + f"PGAI_MAX_CHUNKS_PER_BATCH must be positive, got {max_chunks}. Falling back to provider defaults." + ) + else: + return max_chunks + except ValueError: + logger.warn( + f"PGAI_MAX_CHUNKS_PER_BATCH must be a valid integer, got '{result}'. Falling back to provider defaults." + ) _, custom_llm_provider, _, _ = litellm.get_llm_provider(self.model) # type: ignore match custom_llm_provider: