Repository: run-llama/llama_index
Pull request: see #12234: add support for custom endpoints to MistralAI
# Description

Mistral models can be deployed to Azure AI, which has a different endpoint than the default one. Currently it is not possible to instantiate the `MistralAI` class with a custom endpoint. This PR fixes that by adding an `endpoint` parameter which allows clients to provide a custom endpoint. When not provided, the default endpoint is used as it is now.

Fixes #12234 

## New Package?

- [ ] Yes
- [x] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [x] Yes
- [ ] No

## Type of Change

- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [ ] Added new unit/integration tests
- [ ] Added new notebook (that tests end-to-end)
- [x] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Question]: How to use MistralAI deployed serverless on Azure AI?
### Question Validation

- [X] I have searched both the documentation and discord for an answer.

### Question

I deployed Mistral via Azure AI, how do I use it in LlamaIndex. I tried using the MistralAI LLM in LlamaIndex but I don't see how to configure the endpoint to a custom endpoint.

Hints from issue discussion:
<!-- Greeting -->
Hey @ziodave, nice to see you around here again! Hope you're doing well.

<!-- Answer -->
To configure the MistralAI LLM in LlamaIndex to use a custom endpoint after deploying Mistral via Azure AI, you need to modify the initialization of the `MistralClient` and `MistralAsyncClient` within the `__init__` method of the `MistralAI` class to accept a custom endpoint parameter and pass it to these clients. Here's how you can do it:

1. Add a new parameter `endpoint: Optional[str] = None` to the `__init__` method signature.
2. Use this `endpoint` parameter when initializing `MistralClient` and `MistralAsyncClient`, defaulting to `DEFAULT_MISTRALAI_ENDPOINT` if `None` is provided.

Here's the modified part of the `__init__` method:

```py
def __init__(
        self,
        model: str = DEFAULT_MISTRALAI_MODEL,
        temperature: float = DEFAULT_TEMPERATURE,
        max_tokens: int = DEFAULT_MISTRALAI_MAX_TOKENS,
        timeout: int = 120,
        max_retries: int = 5,
        safe_mode: bool = False,
        random_seed: Optional[int] = None,
        api_key: Optional[str] = None,
        additional_kwargs: Optional[Dict[str, Any]] = None,
        callback_manager: Optional[CallbackManager] = None,
        system_prompt: Optional[str] = None,
        messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
        completion_to_prompt: Optional[Callable[[str], str]] = None,
        pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
        output_parser: Optional[BaseOutputParser] = None,
        endpoint: Optional[str] = None,  # New parameter for custom endpoint
    ) -> None:
        # Existing code to import MistralClient and MistralAsyncClient and check for api_key

        # Use the custom endpoint if provided, otherwise default to DEFAULT_MISTRALAI_ENDPOINT
        endpoint = endpoint or DEFAULT_MISTRALAI_ENDPOINT

        self._client = MistralClient(
            api_key=api_key,
            endpoint=endpoint,  # Use the custom or default endpoint
            timeout=timeout,
            max_retries=max_retries,
        )
        self._aclient = MistralAsyncClient(
            api_key=api_key,
            endpoint=endpoint,  # Use the custom or default endpoint
            timeout=timeout,
            max_retries=max_retries,
        )
        # Rest of the __init__ method...
```

With this modification, you can now pass a custom endpoint URL to the `MistralAI` class upon instantiation, allowing it to communicate with a Mistral deployment on Azure AI or any other platform.

You can find the current structure of the `MistralAI` class and the `DEFAULT_MISTRALAI_ENDPOINT` in the LlamaIndex codebase at the following links:

- [llama-index-legacy/llama_index/legacy/llms/mistral.py](https://github.com/run-llama/llama_index/blob/main/llama-index-legacy/llama_index/legacy/llms/mistral.py)
- [llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py)

Please note that this modification should be done in a fork of the LlamaIndex repository and not directly in the main repository. If you have any further questions, feel free to ask.

For more information, you can refer to the [Dosu](https://dosu.dev) documentation.

## Sources
- [llama-index-legacy/llama_index/legacy/llms/mistral.py](https://github.com/run-llama/llama_index/blob/main/llama-index-legacy/llama_index/legacy/llms/mistral.py)
- [llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py](https://github.com/run-llama/llama_index/blob/main/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py)

<!-- Dosu Comment Footer -->
<details open>
<summary>

#### About Dosu
</summary>

This response is meant to be useful and save you time. It is **not meant to be a precise solution**, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.
</details>

Commit message: see #12234: bump version

Patch to review:
diff --git a/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py b/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py
index 047da4d912..536363ff4c 100644
--- a/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py
+++ b/llama-index-integrations/llms/llama-index-llms-mistralai/llama_index/llms/mistralai/base.py
@@ -143,6 +143,7 @@ def __init__(
         completion_to_prompt: Optional[Callable[[str], str]] = None,
         pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
         output_parser: Optional[BaseOutputParser] = None,
+        endpoint: Optional[str] = None,
     ) -> None:
         additional_kwargs = additional_kwargs or {}
         callback_manager = callback_manager or CallbackManager([])
@@ -155,15 +156,18 @@ def __init__(
                 "You can either pass it in as an argument or set it `MISTRAL_API_KEY`."
             )
 
+        # Use the custom endpoint if provided, otherwise default to DEFAULT_MISTRALAI_ENDPOINT
+        endpoint = endpoint or DEFAULT_MISTRALAI_ENDPOINT
+
         self._client = MistralClient(
             api_key=api_key,
-            endpoint=DEFAULT_MISTRALAI_ENDPOINT,
+            endpoint=endpoint,
             timeout=timeout,
             max_retries=max_retries,
         )
         self._aclient = MistralAsyncClient(
             api_key=api_key,
-            endpoint=DEFAULT_MISTRALAI_ENDPOINT,
+            endpoint=endpoint,
             timeout=timeout,
             max_retries=max_retries,
         )
diff --git a/llama-index-integrations/llms/llama-index-llms-mistralai/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-mistralai/pyproject.toml
index 0564243b38..5b2e270c2f 100644
--- a/llama-index-integrations/llms/llama-index-llms-mistralai/pyproject.toml
+++ b/llama-index-integrations/llms/llama-index-llms-mistralai/pyproject.toml
@@ -27,7 +27,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-mistralai"
 readme = "README.md"
-version = "0.1.9"
+version = "0.1.10"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
diff --git a/llama-index-legacy/llama_index/legacy/llms/mistral.py b/llama-index-legacy/llama_index/legacy/llms/mistral.py
index 4463952ebe..f596051d8b 100644
--- a/llama-index-legacy/llama_index/legacy/llms/mistral.py
+++ b/llama-index-legacy/llama_index/legacy/llms/mistral.py
@@ -91,6 +91,7 @@ def __init__(
         completion_to_prompt: Optional[Callable[[str], str]] = None,
         pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
         output_parser: Optional[BaseOutputParser] = None,
+        endpoint: Optional[str] = None,
     ) -> None:
         try:
             from mistralai.async_client import MistralAsyncClient
@@ -112,15 +113,18 @@ def __init__(
                 "You can either pass it in as an argument or set it `MISTRAL_API_KEY`."
             )
 
+        # Use the custom endpoint if provided, otherwise default to DEFAULT_MISTRALAI_ENDPOINT
+        endpoint = endpoint or DEFAULT_MISTRALAI_ENDPOINT
+
         self._client = MistralClient(
             api_key=api_key,
-            endpoint=DEFAULT_MISTRALAI_ENDPOINT,
+            endpoint=endpoint,
             timeout=timeout,
             max_retries=max_retries,
         )
         self._aclient = MistralAsyncClient(
             api_key=api_key,
-            endpoint=DEFAULT_MISTRALAI_ENDPOINT,
+            endpoint=endpoint,
             timeout=timeout,
             max_retries=max_retries,
         )
diff --git a/llama-index-legacy/pyproject.toml b/llama-index-legacy/pyproject.toml
index 3e671fc07f..1ce0e7d578 100644
--- a/llama-index-legacy/pyproject.toml
+++ b/llama-index-legacy/pyproject.toml
@@ -42,7 +42,7 @@ name = "llama-index-legacy"
 packages = [{include = "llama_index"}]
 readme = "README.md"
 repository = "https://github.com/run-llama/llama_index"
-version = "0.9.48"
+version = "0.9.49"
 
 [tool.poetry.dependencies]
 SQLAlchemy = {extras = ["asyncio"], version = ">=1.4.49"}


Code review comments:
path: llama-index-legacy/llama_index/legacy/llms/mistral.py
text: @user1:
we shouldn't really be updating legacy with new features or version changes

@author:
Done!