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}",