Embeddings API Reference

The maticlib.embeddings module provides a consistent interface for generating text embeddings using various cloud providers.

Base Class

maticlib.embeddings.base.BaseEmbeddings

Bases: ABC

Abstract base class for all embedding models.

embed_documents abstractmethod

embed_documents(texts)

Generates embeddings for a list of document strings.

Parameters:

Name	Type	Description	Default
texts	List[str]	A list of texts to embed.	required

Returns:

Type	Description
EmbedDocumentsResponse	An EmbedDocumentsResponse containing the list of vectors and usage metadata.

Source code in maticlib/embeddings/base.py

@abstractmethod
def embed_documents(self, texts: List[str]) -> EmbedDocumentsResponse:
    """
    Generates embeddings for a list of document strings.

    Args:
        texts: A list of texts to embed.

    Returns:
        An EmbedDocumentsResponse containing the list of vectors and usage metadata.
    """
    pass

embed_query abstractmethod

embed_query(text)

Generates an embedding for a single query string.

Parameters:

Name	Type	Description	Default
text	str	The text to embed.	required

Returns:

Type	Description
EmbedQueryResponse	An EmbedQueryResponse containing the vector and usage metadata.

Source code in maticlib/embeddings/base.py

@abstractmethod
def embed_query(self, text: str) -> EmbedQueryResponse:
    """
    Generates an embedding for a single query string.

    Args:
        text: The text to embed.

    Returns:
        An EmbedQueryResponse containing the vector and usage metadata.
    """
    pass

OpenAI Embeddings

maticlib.embeddings.openai.OpenAIEmbeddings

OpenAIEmbeddings(
    model="text-embedding-3-small",
    api_key=None,
    dimensions=None,
    verbose=True,
)

Bases: BaseEmbeddings

Client for interacting with OpenAI Embedding models.

Parameters:

Name	Type	Description	Default
model	str	The OpenAI embedding model to use. Defaults to "text-embedding-3-small".	'text-embedding-3-small'
api_key	Optional[str]	Your OpenAI API key. Falls back to OPENAI_API_KEY environment variable.	None
dimensions	Optional[int]	The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.	None
verbose	bool	If True, prints status messages to console.	True

Source code in maticlib/embeddings/openai.py

def __init__(
    self,
    model: str = "text-embedding-3-small",
    api_key: Optional[str] = None,
    dimensions: Optional[int] = None,
    verbose: bool = True,
):
    super().__init__()
    api_key = api_key or os.getenv("OPENAI_API_KEY", "")
    api_key = (api_key or "").strip()
    if not api_key:
        raise ValueError(
            "OpenAI API key is missing. Please provide it via the 'api_key' "
            "argument or set the OPENAI_API_KEY environment variable."
        )
    self.api_key = api_key
    self.model = model
    self.dimensions = dimensions
    self.verbose = verbose
    self.base_url = "https://api.openai.com/v1/embeddings"
    self.headers = {
        "Authorization": f"Bearer {self.api_key}",
        "Content-Type": "application/json",
    }

embed_documents

embed_documents(texts)

Embed a list of document strings.

Source code in maticlib/embeddings/openai.py

def embed_documents(self, texts: List[str]) -> EmbedDocumentsResponse:
    """Embed a list of document strings."""
    payload = {
        "model": self.model,
        "input": texts,
    }
    if self.dimensions:
        payload["dimensions"] = self.dimensions

    try:
        response = httpx.post(
            self.base_url, headers=self.headers, json=payload, timeout=60.0
        )
        response.raise_for_status()

        if self.verbose:
            print(f"OpenAI Embeddings Status: {response.status_code}")

        data = response.json()
        usage = data.get("usage", {})
        prompt_tokens = usage.get("prompt_tokens", 0)
        total_tokens = usage.get("total_tokens", prompt_tokens)

        # OpenAI returns data sorted by index in the 'data' list
        embeddings = [
            item["embedding"]
            for item in sorted(data["data"], key=lambda x: x["index"])
        ]

        return EmbedDocumentsResponse(
            vectors=embeddings,
            prompt_tokens=prompt_tokens,
            total_tokens=total_tokens,
            model=data.get("model", self.model),
            raw_response=data,
        )

    except httpx.HTTPStatusError as e:
        if self.verbose:
            print(f"HTTP Error: {e.response.status_code}")
            print(f"Response: {e.response.text}")
        raise
    except Exception:
        raise

embed_query

embed_query(text)

Embed a single query string.

Source code in maticlib/embeddings/openai.py

def embed_query(self, text: str) -> EmbedQueryResponse:
    """Embed a single query string."""
    docs_res = self.embed_documents([text])
    return EmbedQueryResponse(
        vector=docs_res.vectors[0],
        prompt_tokens=docs_res.prompt_tokens,
        total_tokens=docs_res.total_tokens,
        model=docs_res.model,
        raw_response=docs_res.raw_response,
    )

Google GenAI Embeddings

maticlib.embeddings.google.GoogleGenAIEmbeddings

GoogleGenAIEmbeddings(
    model="gemini-embedding-001",
    api_key=None,
    task_type="RETRIEVAL_DOCUMENT",
    verbose=True,
)

Bases: BaseEmbeddings

Client for interacting with Google's Generative AI (Gemini) Embedding models.

Parameters:

Name	Type	Description	Default
model	str	The Gemini embedding model to use. Defaults to "gemini-embedding-001".	'gemini-embedding-001'
api_key	Optional[str]	Your Google AI API key. Falls back to GOOGLE_API_KEY or GEMINI_API_KEY.	None
task_type	str	The type of task the embedding will be used for. Common values: "RETRIEVAL_QUERY", "RETRIEVAL_DOCUMENT", "SEMANTIC_SIMILARITY".	'RETRIEVAL_DOCUMENT'
verbose	bool	If True, prints status messages to console.	True

Source code in maticlib/embeddings/google.py

def __init__(
    self,
    model: str = "gemini-embedding-001",
    api_key: Optional[str] = None,
    task_type: str = "RETRIEVAL_DOCUMENT",
    verbose: bool = True,
):
    super().__init__()
    api_key = (
        api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY") or ""
    )
    api_key = (api_key or "").strip()
    if not api_key:
        raise ValueError(
            "Google Gemini API key is missing. Please provide it via the 'api_key' "
            "argument or set the GOOGLE_API_KEY environment variable."
        )
    self.api_key = api_key
    if not model.startswith("models/"):
        model = f"models/{model}"
    self.model = model
    self.task_type = task_type
    self.verbose = verbose
    self.base_url = "https://generativelanguage.googleapis.com/v1beta"
    self.headers = {
        "x-goog-api-key": self.api_key,
        "Content-Type": "application/json",
    }

embed_documents

embed_documents(texts)

Embed a list of document strings using batchEmbedContents.

Source code in maticlib/embeddings/google.py

def embed_documents(self, texts: List[str]) -> EmbedDocumentsResponse:
    """Embed a list of document strings using batchEmbedContents."""
    url = f"{self.base_url}/{self.model}:batchEmbedContents"

    requests = []
    for text in texts:
        requests.append(
            {
                "model": self.model,
                "content": {"parts": [{"text": text}]},
                "taskType": self.task_type,
            }
        )

    payload = {"requests": requests}

    try:
        response = httpx.post(url, headers=self.headers, json=payload, timeout=60.0)
        response.raise_for_status()

        if self.verbose:
            print(f"Google Batch Embeddings Status: {response.status_code}")

        data = response.json()
        usage = data.get("usageMetadata", {})
        prompt_tokens = usage.get("promptTokenCount", 0)
        vectors = [item["values"] for item in data["embeddings"]]

        return EmbedDocumentsResponse(
            vectors=vectors,
            prompt_tokens=prompt_tokens,
            total_tokens=prompt_tokens,
            model=self.model,
            raw_response=data,
        )

    except Exception as e:
        if self.verbose:
            print(f"Error in Google embed_documents: {e}")
        raise

embed_query

embed_query(text)

Embed a single query string.

Source code in maticlib/embeddings/google.py

def embed_query(self, text: str) -> EmbedQueryResponse:
    """Embed a single query string."""
    url = f"{self.base_url}/{self.model}:embedContent"
    payload = {
        "model": self.model,
        "content": {"parts": [{"text": text}]},
        "taskType": "RETRIEVAL_QUERY",
    }

    try:
        response = httpx.post(url, headers=self.headers, json=payload, timeout=60.0)
        response.raise_for_status()

        if self.verbose:
            print(f"Google Embeddings Status: {response.status_code}")

        data = response.json()
        usage = data.get("usageMetadata", {})
        prompt_tokens = usage.get("promptTokenCount", 0)

        return EmbedQueryResponse(
            vector=data["embedding"]["values"],
            prompt_tokens=prompt_tokens,
            total_tokens=prompt_tokens,
            model=self.model,
            raw_response=data,
        )

    except Exception as e:
        if self.verbose:
            print(f"Error in Google embed_query: {e}")
        raise

Mistral Embeddings

maticlib.embeddings.mistral.MistralEmbeddings

MistralEmbeddings(
    model="mistral-embed", api_key=None, verbose=True
)

Bases: BaseEmbeddings

Client for interacting with Mistral AI Embedding models.

Parameters:

Name	Type	Description	Default
model	str	The Mistral embedding model to use. Defaults to "mistral-embed".	'mistral-embed'
api_key	Optional[str]	Your Mistral API key. Falls back to MISTRAL_API_KEY environment variable.	None
verbose	bool	If True, prints status messages to console.	True

Source code in maticlib/embeddings/mistral.py

def __init__(
    self,
    model: str = "mistral-embed",
    api_key: Optional[str] = None,
    verbose: bool = True,
):
    super().__init__()
    api_key = api_key or os.getenv("MISTRAL_API_KEY", "")
    api_key = (api_key or "").strip()
    if not api_key:
        raise ValueError(
            "Mistral API key is missing. Please provide it via the 'api_key' "
            "argument or set the MISTRAL_API_KEY environment variable."
        )
    self.api_key = api_key
    self.model = model
    self.verbose = verbose
    self.base_url = "https://api.mistral.ai/v1/embeddings"
    self.headers = {
        "Authorization": f"Bearer {self.api_key}",
        "Content-Type": "application/json",
    }

embed_documents

embed_documents(texts)

Embed a list of document strings.

Source code in maticlib/embeddings/mistral.py

def embed_documents(self, texts: List[str]) -> EmbedDocumentsResponse:
    """Embed a list of document strings."""
    payload = {
        "model": self.model,
        "input": texts,
    }

    try:
        response = httpx.post(
            self.base_url, headers=self.headers, json=payload, timeout=60.0
        )
        response.raise_for_status()

        if self.verbose:
            print(f"Mistral Embeddings Status: {response.status_code}")

        data = response.json()
        usage = data.get("usage", {})
        prompt_tokens = usage.get("prompt_tokens", 0)
        total_tokens = usage.get("total_tokens", prompt_tokens)

        # Mistral returns a list of objects with 'embedding' and 'index'
        embeddings = [
            item["embedding"]
            for item in sorted(data["data"], key=lambda x: x["index"])
        ]

        return EmbedDocumentsResponse(
            vectors=embeddings,
            prompt_tokens=prompt_tokens,
            total_tokens=total_tokens,
            model=data.get("model", self.model),
            raw_response=data,
        )

    except Exception as e:
        if self.verbose:
            print(f"Error in Mistral embed_documents: {e}")
        raise

embed_query

embed_query(text)

Embed a single query string.

Source code in maticlib/embeddings/mistral.py

def embed_query(self, text: str) -> EmbedQueryResponse:
    """Embed a single query string."""
    docs_res = self.embed_documents([text])
    return EmbedQueryResponse(
        vector=docs_res.vectors[0],
        prompt_tokens=docs_res.prompt_tokens,
        total_tokens=docs_res.total_tokens,
        model=docs_res.model,
        raw_response=docs_res.raw_response,
    )