From 46bbd559d2270ddb5578b1114b94273f828efe03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=C3=A5ns=20Abrahamsson?= <mans.abrahamsson@gmail.com>
Date: Fri, 24 Jan 2025 14:14:56 +0100
Subject: [PATCH] Added a test-case for the mistralai api

---
 tests/test_llm_api.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/tests/test_llm_api.py b/tests/test_llm_api.py
index 9e2a1d6..405308e 100644
--- a/tests/test_llm_api.py
+++ b/tests/test_llm_api.py
@@ -16,6 +16,25 @@ import sys
 sys.path.append(".")
 
 
+def test_mistral_model():
+    from langchain_core.messages import HumanMessage
+    from src.utils import utils
+
+    llm = utils.get_llm_model(
+        provider="mistral",
+        model_name="mistral-large-latest",
+        temperature=0.8,
+        base_url=os.getenv("MISTRAL_ENDPOINT", ""),
+        api_key=os.getenv("MISTRAL_API_KEY", "")
+    )
+    message = HumanMessage(
+        content=[
+            {"type": "text", "text": "who are you?"}
+        ]
+    )
+    ai_msg = llm.invoke([message])
+    print(ai_msg.content)
+
 def test_openai_model():
     from langchain_core.messages import HumanMessage
     from src.utils import utils
@@ -128,4 +147,5 @@ if __name__ == '__main__':
     # test_gemini_model()
     # test_azure_openai_model()
     # test_deepseek_model()
-    test_ollama_model()
+    # test_ollama_model()
+    test_mistral_model()