From eeefba7ebcab402b0392c246c6dd4909196d9636 Mon Sep 17 00:00:00 2001
From: Arjan Mossel <arjanmossel@gmail.com>
Date: Mon, 25 Nov 2024 18:00:22 +0100
Subject: [PATCH] Add refresh command to fetch available models via API

---
 llm_venice.py | 44 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/llm_venice.py b/llm_venice.py
index d61cdf5..34d4746 100644
--- a/llm_venice.py
+++ b/llm_venice.py
@@ -1,3 +1,7 @@
+import json
+
+import click
+import httpx
 import llm
 from llm.default_plugins.openai_models import Chat
 
@@ -19,13 +23,51 @@ def __str__(self):
         return f"Venice Chat: {self.model_id}"
 
 
+@llm.hookimpl
+def register_commands(cli):
+    @cli.group(name="venice")
+    def venice():
+        "llm-venice plugin commands"
+
+    @venice.command(name="refresh")
+    def refresh():
+        "Refresh the list of models from the Venice API"
+        key = llm.get_key("", "venice", "LLM_VENICE_KEY")
+        if not key:
+            raise click.ClickException("No key found for Venice")
+        headers = {"Authorization": f"Bearer {key}"}
+        response = httpx.get(
+            "https://api.venice.ai/api/v1/models", headers=headers
+        )
+        response.raise_for_status()
+        models = response.json()["data"]
+        text_models = [
+            model["id"]
+            for model in models
+            if model.get("type", {})
+            == "text"
+        ]
+        if not text_models:
+            raise click.ClickException("No text generation models found")
+        path = llm.user_dir() / "llm-venice.json"
+        path.write_text(json.dumps(text_models, indent=4))
+        click.echo(f"{len(text_models)} models saved to {path}", err=True)
+        click.echo(json.dumps(text_models, indent=4))
+
+
 @llm.hookimpl
 def register_models(register):
     key = llm.get_key("", "venice", "LLM_VENICE_KEY")
     if not key:
         return
 
-    for model_id in MODELS:
+    path = llm.user_dir() / "llm-venice.json"
+    if path.exists():
+        model_ids = json.loads(path.read_text())
+    else:
+        model_ids = MODELS
+
+    for model_id in model_ids:
         register(
             VeniceChat(
                 model_id=f"venice/{model_id}",