add missing file ollama.py

ricopinazo · May 8, 2024 · 6e9ae2b · 6e9ae2b
1 parent 3f382cb
commit 6e9ae2b
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/comai/ollama.py 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