Skip to content

Commit

Permalink
Add refresh command to fetch available models via API
Browse files Browse the repository at this point in the history
  • Loading branch information
ar-jan committed Nov 25, 2024
1 parent ea18c6b commit eeefba7
Showing 1 changed file with 43 additions and 1 deletion.
44 changes: 43 additions & 1 deletion llm_venice.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
import json

import click
import httpx
import llm
from llm.default_plugins.openai_models import Chat

Expand All @@ -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}",
Expand Down

0 comments on commit eeefba7

Please sign in to comment.