Skip to content

Rebase PR #898: Add gemini provider support and PGAI_MAX_CHUNKS_PER_BATCH env var #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Copilot wants to merge 3 commits into
base: main
Choose a base branch
from
Draft

Rebase PR #898: Add gemini provider support and PGAI_MAX_CHUNKS_PER_BATCH env var #2

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
112f82e
Initial plan
Copilot Dec 5, 2025
8727b80
feat: rebase patch-1 changes on main - add gemini support and env var…
Copilot Dec 5, 2025
4faf536
fix: add validation for PGAI_MAX_CHUNKS_PER_BATCH env var
Copilot Dec 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions projects/pgai/pgai/vectorizer/embedders/litellm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from collections.abc import AsyncGenerator, Callable
import os
from typing import Any, Literal

from pydantic import BaseModel
Expand Down Expand Up @@ -61,6 +62,20 @@ 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"):
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:
case "cohere":
Expand All @@ -71,6 +86,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":
Expand All @@ -80,10 +97,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:
Expand Down