dllmforge.rag_embedding

This module provides embedding functionality for RAG (Retrieval-Augmented Generation) pipelines. It can be used to 1) vectorize document chunks, and 2) vectorize user queries. The module uses Azure OpenAI embeddings model as an example of using hosted embedding APIs. Note you need Azure OpenAI service and a deployed embedding model on Azure to use this module.

Classes

AzureOpenAIEmbeddingModel([model, api_base, ...])

Class for embedding queries and document chunks using Azure OpenAI Embeddings.

class dllmforge.rag_embedding.AzureOpenAIEmbeddingModel(model: str = 'text-embedding-3-large', api_base: str = None, deployment_name_embeddings: str = None, api_key: str = None, api_version: str = None)

Class for embedding queries and document chunks using Azure OpenAI Embeddings.

Initialize the embedding model using provided arguments or environment variables for Azure OpenAI. :param model: Name of the embedding model to use :param api_base: Azure OpenAI API base URL :param deployment_name_embeddings: Azure OpenAI deployment name for embeddings :param api_key: Azure OpenAI API key :param api_version: Azure OpenAI API version

__init__(model: str = 'text-embedding-3-large', api_base: str = None, deployment_name_embeddings: str = None, api_key: str = None, api_version: str = None)

Initialize the embedding model using provided arguments or environment variables for Azure OpenAI. :param model: Name of the embedding model to use :param api_base: Azure OpenAI API base URL :param deployment_name_embeddings: Azure OpenAI deployment name for embeddings :param api_key: Azure OpenAI API key :param api_version: Azure OpenAI API version

static validate_embedding(embedding: List[float]) → bool

Validate that the embedding is not empty.

static encode_filename(filename: str) → str

Encode filename to be safe for Azure Cognitive Search document keys.

embed(query_or_chunks: str | List[Dict[str, Any]]) → List[float] | List[Dict[str, Any]]

Embed a single query string or a list of document chunks. :param query_or_chunks: A string (query) or a list of dictionaries (document chunks)

Each dictionary should have keys: “text”, “file_name”, “page_number”

Returns:

list of floats (embedding vector) For document chunks: list of dictionaries with keys: “chunk_id”, “chunk”, “page_number”,

”file_name”, “text_vector”

Return type:

For a string query