rag_system.openrouter_embeddings module

Embeddings client for RAG system.

Uses the native Google Gemini API only, via the shared key pool in gemini_embed_pool. Provides both async (OpenRouterEmbeddings) and synchronous (SyncOpenRouterEmbeddings, ChromaDB-compatible) interfaces.

class rag_system.openrouter_embeddings.OpenRouterEmbeddings(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True)[source]

Bases: object

Async embeddings client using Gemini API via the shared key pool.

Parameters:

  • api_key (str | None)

  • model (str)

  • dimensions (int | None)

  • timeout (float)

  • gemini_api_key (str | None)

  • gemini_only (bool)

DEFAULT_MODEL = 'google/gemini-embedding-001'
MAX_BATCH_SIZE = 50
MAX_BATCH_CHARS = 50000
__init__(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True)[source]

Initialize the instance.

Parameters:

  • api_key (Optional[str]) – Unused; kept for backward compatibility.

  • model (str) – The model value.

  • dimensions (Optional[int]) – The dimensions value.

  • timeout (float) – Maximum wait time in seconds.

  • gemini_api_key (Optional[str]) – Unused; pool is used instead.

  • gemini_only (bool) – Always True; embeddings use Gemini API only.

async embed_text(text)[source]

Embed text.

Parameters:

text (str) – Text content.

Returns:

The result.

Return type:

ndarray

async embed_texts(texts)[source]

Embed texts.

Parameters:

texts (List[str]) – The texts value.

Returns:

The result.

Return type:

List[ndarray]

async embed_text_for_search(text, task_type='QUESTION_ANSWERING')[source]

Embed a single text using the Gemini API only, with a task type.

Intended for pre-computing a query embedding before passing it to FileRAGManager.search(query_embedding=...). Retries on transient errors with exponential back-off.

Return type:

List[float]

Parameters:
async close()[source]

Close.

async __aenter__()[source]

Internal helper: aenter .

async __aexit__(exc_type, exc_val, exc_tb)[source]

Internal helper: aexit .

Parameters:

  • exc_type – The exc type value.

  • exc_val – The exc val value.

  • exc_tb – The exc tb value.

class rag_system.openrouter_embeddings.SyncOpenRouterEmbeddings(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True, document_task_type=None, query_task_type=None)[source]

Bases: object

Synchronous wrapper used by ChromaDB’s embedding function interface.

Uses Gemini API via the shared key pool. Batches are dispatched concurrently via a ThreadPoolExecutor when there are multiple batches.

Parameters:

  • api_key (str | None)

  • model (str)

  • dimensions (int | None)

  • timeout (float)

  • gemini_api_key (str | None)

  • gemini_only (bool)

  • document_task_type (str | None)

  • query_task_type (str | None)

MAX_BATCH_SIZE = 50
MAX_BATCH_CHARS = 50000
MAX_EMBED_WORKERS = 20
__init__(api_key=None, model='google/gemini-embedding-001', dimensions=None, timeout=30.0, gemini_api_key=None, gemini_only=True, document_task_type=None, query_task_type=None)[source]

Initialize the instance.

Parameters:

  • api_key (Optional[str]) – Unused; kept for backward compatibility.

  • model (str) – The model value.

  • dimensions (Optional[int]) – The dimensions value.

  • timeout (float) – Maximum wait time in seconds.

  • gemini_api_key (Optional[str]) – Unused; pool is used instead.

  • gemini_only (bool) – Unused; always Gemini API.

  • document_task_type (Optional[str]) – Optional Gemini taskType for corpus (e.g. RETRIEVAL_DOCUMENT); used by embed_documents.

  • query_task_type (Optional[str]) – Optional Gemini taskType for queries (e.g. RETRIEVAL_QUERY); used by embed_query.

name()[source]

Name.

Returns:

Result string.

Return type:

str

dimension()[source]

Dimension.

Returns:

The result.

Return type:

int

__call__(input)[source]

ChromaDB EmbeddingFunction interface (legacy).

Uses document_task_type when set (same as embed_documents()).

Return type:

List[List[float]]

Parameters:

input (List[str])

embed_documents(input)[source]

Embed documents (ChromaDB interface for upsert).

Return type:

List[List[float]]

Parameters:

input (List[str])

embed_query(input)[source]

Embed query texts (ChromaDB interface for query).

Return type:

List[List[float]]

Parameters:

input (List[str])