Update mistralai_.py

docs/examples/kickstart_snippets/mistralai_.py

@@ -9,54 +9,69 @@
 - Uses `yield` to continuously concatenate the parts of the response.
 """
 
+import os
+
 import panel as pn
-from mistralai.async_client import MistralAsyncClient
+from mistralai import Mistral, UserMessage
 
 pn.extension()
 
 
 def update_api_key(api_key):
-    # use api_key_input.value if set, otherwise use MISTRAL_API_KEY
+    # Use the provided api_key or default to the environment variable
     pn.state.cache["aclient"] = (
-        MistralAsyncClient(api_key=api_key) if api_key else MistralAsyncClient()
+        Mistral(api_key=api_key)
+        if api_key
+        else Mistral(api_key=os.getenv("MISTRAL_API_KEY", ""))
     )
 
 
 async def callback(contents: str, user: str, instance: pn.chat.ChatInterface):
-    # memory is a list of messages
+    # memory is a list of serialized messages
     messages = instance.serialize()
 
-    response = pn.state.cache["aclient"].chat_stream(
+    # Convert serialized messages into UserMessage format
+    formatted_messages = [UserMessage(content=msg["content"]) for msg in messages]
+
+    response = await pn.state.cache["aclient"].chat.stream_async(
         model="mistral-small",
-        messages=messages,
+        messages=formatted_messages,
     )
 
     message = ""
     async for chunk in response:
-        part = chunk.choices[0].delta.content
+        part = chunk.data.choices[0].delta.content
         if part is not None:
             message += part
             yield message
 
 
+# Input widget for the API key
 api_key_input = pn.widgets.PasswordInput(
     placeholder="Uses $MISTRAL_API_KEY if not set",
     sizing_mode="stretch_width",
     styles={"color": "black"},
 )
+
+# Bind the API key input to the update function
 pn.bind(update_api_key, api_key_input, watch=True)
 api_key_input.param.trigger("value")
 
+# Define the Chat Interface with callback
 chat_interface = pn.chat.ChatInterface(
     callback=callback,
     callback_user="MistralAI",
     help_text="Send a message to get a reply from MistralAI!",
     callback_exception="verbose",
 )
+
+# Template with the chat interface
 template = pn.template.FastListTemplate(
     title="MistralAI Small",
     header_background="#FF7000",
     main=[chat_interface],
     header=[api_key_input],
 )
+
+# Serve the template
 template.servable()