AI Blogen - WIP - 0.0.00000.001

lib/gpt_providers/mistral_chat_completion.py

@@ -0,0 +1,40 @@
+import os
+import logging
+from pathlib import Path
+
+from mistralai.client import MistralClient
+from mistralai.models.chat_completion import ChatMessage
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s-%(levelname)s-%(module)s-%(lineno)d-%(message)s')
+from dotenv import load_dotenv
+load_dotenv(Path('../../.env'))
+
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)  # for exponential backoff
+
+
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
+def mistral_text_response(prompt):
+    """ Common function to get text response from minstral. """
+    api_key = os.environ["MISTRAL_API_KEY"]
+    model = "mistral-medium"
+
+    client = MistralClient(api_key=api_key)
+
+    messages = [
+        ChatMessage(role="user", content=prompt)
+    ]
+
+    # No streaming
+    chat_response = client.chat(
+        model=model,
+        messages=messages,
+    )
+    print(chat_response)
+
+    # With streaming
+    for chunk in client.chat_stream(model=model, messages=messages):
+        print(chunk)