Skip to content

Add Mistral AI batch processing example to MistralEmbeddingFunction #5597

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

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add Mistral AI batch processing example to MistralEmbeddingFunction #5597

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c4abe4a
Add Mistral AI batch processing example to MistralEmbeddingFunction
Oct 13, 2025
1288d62
Update docstrings
Oct 13, 2025
f22c63b
Add propel-code-bot
Oct 13, 2025
fd01385
Merge branch 'chroma-core:main' into Handle-Mistral-AI-text-embedding…
Gwenn-LR Oct 14, 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: 16 additions & 6 deletions chromadb/utils/embedding_functions/mistral_embedding_function.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,15 @@ def __init__(
self,
model: str,
api_key_env_var: str = "MISTRAL_API_KEY",
max_chunk_size: int = 50
):
"""
Initialize the MistralEmbeddingFunction.

Args:
model (str): The name of the model to use for text embeddings.
api_key_env_var (str): The environment variable name for the Mistral API key.
max_chunk_size (int): The maximal size of chunk to embed. Default to 50.
"""
try:
from mistralai import Mistral
Expand All @@ -30,23 +32,29 @@ def __init__(
if not self.api_key:
raise ValueError(f"The {api_key_env_var} environment variable is not set.")
self.client = Mistral(api_key=self.api_key)
self.max_chunk_size = max_chunk_size

def __call__(self, input: Documents) -> Embeddings:
"""
Get the embeddings for a list of texts.

Note:
Mistral AI API will provide auto-chunking in the future (https://docs.mistral.ai/capabilities/embeddings/text_embeddings/#batch-processing)

Args:
input (Documents): A list of texts to get embeddings for.
"""
if not all(isinstance(item, str) for item in input):
raise ValueError("Mistral only supports text documents, not images")
output = self.client.embeddings.create(
model=self.model,
inputs=input,

chunks = (input[x : x + self.max_chunk_size] for x in range(0, len(input), self.max_chunk_size))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Empty input handling is missing. If input is an empty list, chunks will be an empty generator and the final list comprehension will return an empty list, which may be correct behavior. However, if this is intentional, consider adding an explicit check for clarity:

Suggested change
chunks = (input[x : x + self.max_chunk_size] for x in range(0, len(input), self.max_chunk_size))
if not input:
return []
chunks = (input[x : x + self.max_chunk_size] for x in range(0, len(input), self.max_chunk_size))

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents 
[**BestPractice**]

Empty input handling is missing. If `input` is an empty list, `chunks` will be an empty generator and the final list comprehension will return an empty list, which may be correct behavior. However, if this is intentional, consider adding an explicit check for clarity:

```suggestion
        if not input:
            return []
            
        chunks = (input[x : x + self.max_chunk_size] for x in range(0, len(input), self.max_chunk_size))
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

File: chromadb/utils/embedding_functions/mistral_embedding_function.py
Line: 50

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I haven't added those changes since the test above should raise an error before.

embeddings_response = (
self.client.embeddings.create(model=self.model, inputs=c) for c in chunks
Comment on lines +51 to +52
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[BestPractice]

Missing Error Handling: API calls to self.client.embeddings.create() can fail due to network issues, API limits, authentication errors, or the specific "Too many tokens" error mentioned in the PR description. Without error handling, these failures will propagate as unhandled exceptions.

Add error handling:

for chunk in chunks:
    try:
        response = self.client.embeddings.create(model=self.model, inputs=chunk)
        all_embeddings.extend([np.array(d.embedding) for d in response.data])
    except Exception as e:
        raise RuntimeError(f"Failed to get embeddings for chunk: {e}") from e
Context for Agents 
[**BestPractice**]

**Missing Error Handling**: API calls to `self.client.embeddings.create()` can fail due to network issues, API limits, authentication errors, or the specific "Too many tokens" error mentioned in the PR description. Without error handling, these failures will propagate as unhandled exceptions.

**Add error handling:**
```python
for chunk in chunks:
    try:
        response = self.client.embeddings.create(model=self.model, inputs=chunk)
        all_embeddings.extend([np.array(d.embedding) for d in response.data])
    except Exception as e:
        raise RuntimeError(f"Failed to get embeddings for chunk: {e}") from e
```

File: chromadb/utils/embedding_functions/mistral_embedding_function.py
Line: 52

)


# Extract embeddings from the response
return [np.array(data.embedding) for data in output.data]
return [np.array(d.embedding) for e in embeddings_response for d in e.data]

@staticmethod
def name() -> str:
Expand All @@ -62,15 +70,17 @@ def supported_spaces(self) -> List[Space]:
def build_from_config(config: Dict[str, Any]) -> "EmbeddingFunction[Documents]":
model = config.get("model")
api_key_env_var = config.get("api_key_env_var")
max_chunk_size = config.get("max_chunk_size")

if model is None or api_key_env_var is None:
if model is None or api_key_env_var is None or max_chunk_size is None:
assert False, "This code should not be reached" # this is for type checking
return MistralEmbeddingFunction(model=model, api_key_env_var=api_key_env_var)
return MistralEmbeddingFunction(model=model, api_key_env_var=api_key_env_var, max_chunk_size=max_chunk_size)

def get_config(self) -> Dict[str, Any]:
return {
"model": self.model,
"api_key_env_var": self.api_key_env_var,
"max_chunk_size": self.max_chunk_size
}

def validate_config_update(
Expand Down