embeddings.py — langchain Source File

Source Code

import asyncio
import logging
import warnings
from collections.abc import Callable, Iterable

import httpx
from httpx import Response
from langchain_core.embeddings import Embeddings
from langchain_core.utils import (
    secret_from_env,
)
from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    SecretStr,
    model_validator,
)
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_fixed
from tokenizers import Tokenizer  # type: ignore[import]
from typing_extensions import Self

logger = logging.getLogger(__name__)

MAX_TOKENS = 16_000
"""A batching parameter for the Mistral API. This is NOT the maximum number of tokens
accepted by the embedding model for each document/chunk, but rather the maximum number
of tokens that can be sent in a single request to the Mistral API (across multiple
documents/chunks)"""


class DummyTokenizer:
    """Dummy tokenizer for when tokenizer cannot be accessed (e.g., via Huggingface)."""

    @staticmethod
    def encode_batch(texts: list[str]) -> list[list[str]]:
        return [list(text) for text in texts]


class MistralAIEmbeddings(BaseModel, Embeddings):
    """MistralAI embedding model integration.

    Setup:
        Install `langchain_mistralai` and set environment variable
        `MISTRAL_API_KEY`.

        ```bash
        pip install -U langchain_mistralai
        export MISTRAL_API_KEY="your-api-key"
        ```

    Key init args — completion params:
        model:
            Name of `MistralAI` model to use.

    Key init args — client params:
        api_key:
            The API key for the MistralAI API. If not provided, it will be read from the
            environment variable `MISTRAL_API_KEY`.
        max_concurrent_requests: int
// ... (270 more lines)