Back to Llama Index

MistralAI

docs/examples/llm/mistralai.ipynb

0.14.216.3 KB
Original Source

<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/mistralai.ipynb" target="_parent"></a>

MistralAI

If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

python
%pip install llama-index-llms-mistralai
python
!pip install llama-index

Call complete with a prompt

python
from llama_index.llms.mistralai import MistralAI

# To customize your API key, do this
# otherwise it will lookup MISTRAL_API_KEY from your env variable
# llm = MistralAI(api_key="<api_key>")

llm = MistralAI(api_key="<replace-with-your-key>")

resp = llm.complete("Paul Graham is ")
python
print(resp)

Call chat with a list of messages

python
from llama_index.core.llms import ChatMessage
from llama_index.llms.mistralai import MistralAI

messages = [
    ChatMessage(role="system", content="You are CEO of MistralAI."),
    ChatMessage(role="user", content="Tell me the story about La plateforme"),
]
resp = MistralAI().chat(messages)
python
print(resp)

Call with random_seed

python
from llama_index.core.llms import ChatMessage
from llama_index.llms.mistralai import MistralAI

messages = [
    ChatMessage(role="system", content="You are CEO of MistralAI."),
    ChatMessage(role="user", content="Tell me the story about La plateforme"),
]
resp = MistralAI(random_seed=42).chat(messages)
python
print(resp)

Streaming

Using stream_complete endpoint

python
from llama_index.llms.mistralai import MistralAI

llm = MistralAI()
resp = llm.stream_complete("Paul Graham is ")
python
for r in resp:
    print(r.delta, end="")
python
from llama_index.llms.mistralai import MistralAI
from llama_index.core.llms import ChatMessage

llm = MistralAI()
messages = [
    ChatMessage(role="system", content="You are CEO of MistralAI."),
    ChatMessage(role="user", content="Tell me the story about La plateforme"),
]
resp = llm.stream_chat(messages)
python
for r in resp:
    print(r.delta, end="")

Configure Model

python
from llama_index.llms.mistralai import MistralAI

llm = MistralAI(model="mistral-medium")
python
resp = llm.stream_complete("Paul Graham is ")
python
for r in resp:
    print(r.delta, end="")

Function Calling

mistral-large supports native function calling. There's a seamless integration with LlamaIndex tools, through the predict_and_call function on the llm.

This allows the user to attach any tools and let the LLM decide which tools to call (if any).

If you wish to perform tool calling as part of an agentic loop, check out our agent guides instead.

NOTE: If you use another Mistral model, we will use a ReAct prompt to attempt to call the function. Your mileage may vary.

python
from llama_index.llms.mistralai import MistralAI
from llama_index.core.tools import FunctionTool


def multiply(a: int, b: int) -> int:
    """Multiple two integers and returns the result integer"""
    return a * b


def mystery(a: int, b: int) -> int:
    """Mystery function on two integers."""
    return a * b + a + b


mystery_tool = FunctionTool.from_defaults(fn=mystery)
multiply_tool = FunctionTool.from_defaults(fn=multiply)

llm = MistralAI(model="mistral-large-latest")
python
response = llm.predict_and_call(
    [mystery_tool, multiply_tool],
    user_msg="What happens if I run the mystery function on 5 and 7",
)
python
print(str(response))
python
response = llm.predict_and_call(
    [mystery_tool, multiply_tool],
    user_msg=(
        """What happens if I run the mystery function on the following pairs of numbers? Generate a separate result for each row:
- 1 and 2
- 8 and 4
- 100 and 20 \
"""
    ),
    allow_parallel_tool_calls=True,
)
python
print(str(response))
python
for s in response.sources:
    print(f"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}")

You get the same result if you use the async variant (it will be faster since we do asyncio.gather under the hood).

python
response = await llm.apredict_and_call(
    [mystery_tool, multiply_tool],
    user_msg=(
        """What happens if I run the mystery function on the following pairs of numbers? Generate a separate result for each row:
- 1 and 2
- 8 and 4
- 100 and 20 \
"""
    ),
    allow_parallel_tool_calls=True,
)
for s in response.sources:
    print(f"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}")

Structured Prediction

An important use case for function calling is extracting structured objects. LlamaIndex provides an intuitive interface for converting any LLM into a structured LLM - simply define the target Pydantic class (can be nested), and given a prompt, we extract out the desired object.

python
from llama_index.llms.mistralai import MistralAI
from llama_index.core.prompts import PromptTemplate
from llama_index.core.bridge.pydantic import BaseModel


class Restaurant(BaseModel):
    """A restaurant with name, city, and cuisine."""

    name: str
    city: str
    cuisine: str


llm = MistralAI(model="mistral-large-latest")
prompt_tmpl = PromptTemplate(
    "Generate a restaurant in a given city {city_name}"
)

# Option 1: Use `as_structured_llm`
restaurant_obj = (
    llm.as_structured_llm(Restaurant)
    .complete(prompt_tmpl.format(city_name="Miami"))
    .raw
)
# Option 2: Use `structured_predict`
# restaurant_obj = llm.structured_predict(Restaurant, prompt_tmpl, city_name="Miami")
python
restaurant_obj

Structured Prediction with Streaming

Any LLM wrapped with as_structured_llm supports streaming through stream_chat.

python
from llama_index.core.llms import ChatMessage
from IPython.display import clear_output
from pprint import pprint

input_msg = ChatMessage.from_str("Generate a restaurant in Miami")

sllm = llm.as_structured_llm(Restaurant)
stream_output = sllm.stream_chat([input_msg])
for partial_output in stream_output:
    clear_output(wait=True)
    pprint(partial_output.raw.dict())
    restaurant_obj = partial_output.raw

restaurant_obj

Async

python
from llama_index.llms.mistralai import MistralAI

llm = MistralAI()
resp = await llm.acomplete("Paul Graham is ")
python
print(resp)