diff --git a/nodes/utilities.py b/nodes/utilities.py index b772ab1..b106b4d 100644 --- a/nodes/utilities.py +++ b/nodes/utilities.py @@ -9,6 +9,29 @@ from jinja2 import Template from PIL import Image, ImageOps, ImageSequence +def get_models(engine, base_url, port) -> list[str]: + if engine == "ollama": + api_url = f"http://{base_url}:{port}/api/tags" + try: + response = requests.get(api_url) + response.raise_for_status() + models = [model["name"] for model in response.json().get("models", [])] + return models + except Exception as e: + print(f"Failed to fetch models from Ollama: {e}") + return [] + elif engine == "lmstudio": + api_url = f"http://{base_url}:{port}/v1/models" + try: + response = requests.get(api_url) + response.raise_for_status() + models = [model["id"] for model in response.json().get("data", [])] + return models + except Exception as e: + print(f"Failed to fetch models from LM Studio: {e}") + return [] + + def get_lmstudio_models(port="1234") -> list[str]: url = f"http://127.0.0.1:{port}/v1/models" diff --git a/requirements.txt b/requirements.txt index 0ff65ff..f73451e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ griptape[drivers-prompt-cohere,drivers-prompt-anthropic,drivers-prompt-huggingface,drivers-prompt-amazon-bedrock,drivers-prompt-amazon-sagemaker,drivers-prompt-google,drivers-prompt-ollama,drivers-embedding-amazon-bedrock,drivers-embedding-amazon-sagemaker,drivers-embedding-huggingface,drivers-embedding-voyageai,drivers-embedding-google,drivers-embedding-cohere,drivers-vector-pinecone,drivers-vector-marqo,drivers-vector-mongodb,drivers-vector-redis,drivers-vector-opensearch,drivers-vector-postgresql,drivers-vector-qdrant,drivers-web-scraper-trafilatura,drivers-web-scraper-markdownify,drivers-web-search-duckduckgo,loaders-dataframe,loaders-pdf,loaders-image,loaders-email,loaders-audio]==0.28.2 python-dotenv +elevenlabs \ No newline at end of file