Support for Mistral Large and Mistral Nemo on Google Vertex AI

[dotrunghieu96]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Google Model Garden now supported Mistral Large model family, including Mistral Nemo and Codestral. The LLM ChatVertexAI currently doesn't work with these model.
Please take a look at these, a package similar to langchain-anthropic would be awesome.

Motivation

Our application depends heavily on langchain, and integrating new LLMs without it being supported on Langchain is troublesome.

Proposal (If applicable)

No response

[calls9-amirbraham]

@dotrunghieu96 Any solutions for this ?

[dotrunghieu96]

@calls9-amirbraham checkout my implemetation below

from __future__ import annotations

import logging
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Type, cast

import google.auth
import httpx
from google.auth.transport.requests import Request
from httpx_sse import aconnect_sse, connect_sse
from langchain_core.callbacks import (
    AsyncCallbackManagerForLLMRun,
    CallbackManagerForLLMRun,
)
from langchain_core.language_models.chat_models import agenerate_from_stream
from langchain_core.messages import AIMessageChunk, BaseMessage, BaseMessageChunk
from langchain_core.outputs import ChatGenerationChunk, ChatResult
from langchain_mistralai.chat_models import (
    ChatMistralAI,
    _aiter_sse,
    _araise_on_error,
    _convert_chunk_to_message_chunk,
    _create_retry_decorator,
    _raise_on_error,
)
from pydantic import Field, model_validator

logger = logging.getLogger(__name__)


def get_gcloud_credentials():
    credentials, project_id = google.auth.default(
        scopes=["https://www.googleapis.com/auth/cloud-platform"]
    )
    credentials.refresh(Request())
    return credentials.token, project_id


def build_endpoint_url(
    region: str,
    project_id: str,
):
    base_url = f"https://{region}-aiplatform.googleapis.com/v1/"
    project_fragment = f"projects/{project_id}"
    location_fragment = f"locations/{region}"
    model_fragment = f"publishers/mistralai/models/"
    url = f"{base_url}{'/'.join([project_fragment, location_fragment, model_fragment])}"
    return url


async def acompletion_with_retry(
    llm: ChatMistralVertexAI,
    run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
    **kwargs: Any,
) -> Any:
    """Use tenacity to retry the async completion call."""
    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)

    @retry_decorator
    async def _completion_with_retry(**kwargs: Any) -> Any:
        kwargs.pop("safe_prompt", None)
        if "stream" not in kwargs:
            kwargs["stream"] = False
        stream = kwargs["stream"]
        if stream:
            logger.debug("connect event source")
            event_source = aconnect_sse(
                llm.async_client,
                "POST",
                url=f"{llm.model}@{llm.model_version}:streamRawPredict",
                json=kwargs,
            )
            logger.debug("connected event source")
            return _aiter_sse(event_source)
        else:
            response = await llm.async_client.post(
                url=f"{llm.model}@{llm.model_version}:streamRawPredict", json=kwargs
            )
            await _araise_on_error(response)
            return response.json()

    return await _completion_with_retry(**kwargs)


class ChatMistralVertexAI(ChatMistralAI):
    model: str = Field(default="mistral-large", alias="model_name")
    model_version: Optional[str] = None
    location: str = Field(default="europe-west4")
    project_id: Optional[str] = None
    endpoint: Optional[str] = None

    @model_validator(mode="before")
    @classmethod
    def validate_input_environment(cls, values: Dict) -> Dict:
        """Validate api key, python package exists, temperature, and top_p."""

        if "@" in values["model"]:
            splited_model_name = values["model"].split("@")
            values["model"] = splited_model_name[0]
            values["model_version"] = splited_model_name[1]

        gcloud_access_token, project_id = get_gcloud_credentials()

        values["mistral_api_key"] = gcloud_access_token
        api_key_str = gcloud_access_token
        values["project_id"] = project_id
        if not values.get("location"):
            values["location"] = "europe-west4"
        values["endpoint"] = build_endpoint_url(
            values["location"], values["project_id"]
        )
        # todo: handle retries
        if not values.get("client"):
            values["client"] = httpx.Client(
                base_url=values["endpoint"],
                headers={
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                    "Authorization": f"Bearer {api_key_str}",
                },
                timeout=values.get("timeout", 120),
            )
        # todo: handle retries and max_concurrency
        if not values.get("async_client"):
            values["async_client"] = httpx.AsyncClient(
                base_url=values["endpoint"],
                headers={
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                    "Authorization": f"Bearer {api_key_str}",
                },
                timeout=values.get("timeout", 120),
            )

        if values["temperature"] is not None and not 0 <= values["temperature"] <= 1:
            raise ValueError("temperature must be in the range [0.0, 1.0]")

        if values["top_p"] is not None and not 0 <= values["top_p"] <= 1:
            raise ValueError("top_p must be in the range [0.0, 1.0]")

        return values

    def completion_with_retry(
        self, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any
    ) -> Any:
        """Use tenacity to retry the completion call."""
        # retry_decorator = _create_retry_decorator(self, run_manager=run_manager)

        # @retry_decorator
        def _completion_with_retry(**kwargs: Any) -> Any:
            kwargs.pop("safe_prompt", None)
            if "stream" not in kwargs:
                kwargs["stream"] = False
            stream = kwargs["stream"]
            if stream:

                def iter_sse() -> Iterator[Dict]:
                    with connect_sse(
                        self.client,
                        "POST",
                        f"/{self.model}:rawPredict",
                        json=kwargs,
                    ) as event_source:
                        _raise_on_error(event_source.response)
                        for event in event_source.iter_sse():
                            if event.data == "[DONE]":
                                return
                            yield event.json()

                return iter_sse()
            else:
                response = self.client.post(
                    url=f"/{self.model}:rawPredict", json=kwargs
                )
                _raise_on_error(response)
                return response.json()

        rtn = _completion_with_retry(**kwargs)
        return rtn

    async def _astream(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> AsyncIterator[ChatGenerationChunk]:
        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {**params, **kwargs, "stream": True}

        default_chunk_class: Type[BaseMessageChunk] = AIMessageChunk
        async for chunk in await acompletion_with_retry(
            self, messages=message_dicts, run_manager=run_manager, **params
        ):
            if len(chunk["choices"]) == 0:
                continue
            new_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)
            # make future chunks same type as first chunk
            default_chunk_class = new_chunk.__class__
            gen_chunk = ChatGenerationChunk(message=new_chunk)
            if run_manager:
                await run_manager.on_llm_new_token(
                    token=cast(str, new_chunk.content), chunk=gen_chunk
                )
            yield gen_chunk

    async def _agenerate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
        stream: Optional[bool] = None,
        **kwargs: Any,
    ) -> ChatResult:
        should_stream = stream if stream is not None else self.streaming
        if should_stream:
            stream_iter = self._astream(
                messages=messages, stop=stop, run_manager=run_manager, **kwargs
            )
            return await agenerate_from_stream(stream_iter)

        message_dicts, params = self._create_message_dicts(messages, stop)
        params = {**params, **kwargs}
        response = await acompletion_with_retry(
            self, messages=message_dicts, run_manager=run_manager, **params
        )
        return self._create_chat_result(response)

[dotrunghieu96]

It only works with ainvoke and astream_events though