From 697cf702329a1693790737d4ef37d70956dc2840 Mon Sep 17 00:00:00 2001
From: Jeremy Tuloup <jeremy.tuloup@gmail.com>
Date: Wed, 10 Jul 2024 10:34:38 +0200
Subject: [PATCH] Update the list of models

Co-Authored-By: xiayu <xiayu@kingsoft.com>
---
 .../jupyter_ai_magics/embedding_providers.py           | 10 +++++++++-
 .../jupyter-ai-magics/jupyter_ai_magics/providers.py   |  8 +++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
index 0d8e0df12..dca391da0 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
@@ -71,9 +71,17 @@ class OllamaEmbeddingsProvider(BaseEmbeddingsProvider, OllamaEmbeddings):
     id = "ollama"
     name = "Ollama"
     models = [
+        "gemma",
+        "gemma2",
         "llama2",
-        "mistral:text",
+        "llama3",
+        "llama3:70b",
+        "phi3",
+        "mistral",
         "tinyllama",
+        "qwen2",
+        "qwen2:7b",
+        "qwen2:72b",
     ]
     model_id_key = "model"
 
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
index 2650e2e28..df59354b7 100644
--- a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
+++ b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -728,12 +728,18 @@ class OllamaProvider(BaseProvider, Ollama):
     id = "ollama"
     name = "Ollama"
     model_id_key = "model"
-    # TODO: add more
     models = [
         "gemma",
+        "gemma2",
         "llama2",
+        "llama3",
+        "llama3:70b",
+        "phi3",
         "mistral",
         "tinyllama",
+        "qwen2",
+        "qwen2:7b",
+        "qwen2:72b",
     ]