mistral_chat.py

mistral_chat.py

#!python
# Copyright Ram Rachum 2023, MIT license

# Use this command line script to have a simple interactive chat with the Mistral AI, similar to
# ChatGPT.
#
# First supply your API key like so:
# echo '{"mistral_api_key": "YOUR_API_KEY"}' > ~/.mistral_chat

import os
import json
import pathlib

import click
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

def multi_line_prompt(message: str = '') -> str:
    lines = []
    print(message)
    while True:
        try:
            line = input()
        except EOFError:
            break
        lines.append(line)
    return '\n'.join(lines)
    

@click.command()
def chat_app():
    try:
        api_key = json.loads((pathlib.Path.home() / '.mistral_chat').read_text())['mistral_api_key']
    except (FileNotFoundError, json.JSONDecodeError) as e:
        raise click.UsageError(
            '''To use mistral_chat.py, please run this: \n'''
            '''echo '{"mistral_api_key": "YOUR_API_KEY"}' > ~/.mistral_chat'''
        )
        
    mistral_client = MistralClient(api_key=api_key)
    model = 'mistral-medium'

    click.echo(f'Mistral chat with `{model}`.')
        
    eof_marker = 'Z' if os.name == 'nt' else 'D'
    
    messages = []
    while True:
        user_message = multi_line_prompt(f'You (To finish message: Enter, ^{eof_marker}, Enter): ')
        if user_message.lower() in ['exit', 'quit', '']:
            break
        messages.append(ChatMessage(role='user', content=user_message))
        response = mistral_client.chat(model=model, messages=messages).choices[0].message.content
        messages.append(ChatMessage(role='assistant', content=response))
        click.echo(f'{model}:')
        click.echo(response)

if __name__ == '__main__':
    chat_app()