From 6e9ae2bf7a98bf5acc3fee6e5a23d734afd5864c Mon Sep 17 00:00:00 2001
From: Pedro Rico Pinazo <ricopinazo@gmail.com>
Date: Wed, 8 May 2024 07:02:10 +0100
Subject: [PATCH] add missing file ollama.py

---
 src/comai/ollama.py | 8 ++++++++
 1 file changed, 8 insertions(+)
 create mode 100644 src/comai/ollama.py

diff --git a/src/comai/ollama.py b/src/comai/ollama.py
new file mode 100644
index 0000000..8a1d99a
--- /dev/null
+++ b/src/comai/ollama.py
@@ -0,0 +1,8 @@
+import ollama
+
+
+def get_ollama_model_names() -> list[str]:
+    models = ollama.list()["models"]
+    model_names = [model["name"] for model in models]
+    cleaned_model_names = [name.replace(":latest", "") for name in model_names]
+    return cleaned_model_names