From 6bb70ac1faf3b9b1afaec66b819bb71b0d3a9915 Mon Sep 17 00:00:00 2001
From: Zain Hoda <7146154+zainhoda@users.noreply.github.com>
Date: Tue, 30 Apr 2024 17:26:17 -0400
Subject: [PATCH] add import tests

---
 pyproject.toml        | 3 ++-
 tests/test_imports.py | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index d020f241..aae6d9fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,7 +32,7 @@ bigquery = ["google-cloud-bigquery"]
 snowflake = ["snowflake-connector-python"]
 duckdb = ["duckdb"]
 google = ["google-generativeai", "google-cloud-aiplatform"]
-all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl"]
+all = ["psycopg2-binary", "db-dtypes", "PyMySQL", "google-cloud-bigquery", "snowflake-connector-python", "duckdb", "openai", "mistralai", "chromadb", "anthropic", "zhipuai", "marqo", "google-generativeai", "google-cloud-aiplatform", "qdrant-client", "fastembed", "ollama", "httpx", "opensearch-py", "opensearch-dsl", "transformers"]
 test = ["tox"]
 chromadb = ["chromadb"]
 openai = ["openai"]
@@ -45,3 +45,4 @@ ollama = ["ollama", "httpx"]
 qdrant = ["qdrant-client"]
 vllm = ["vllm"]
 opensearch = ["opensearch-py", "opensearch-dsl"]
+hf = ["transformers"]
diff --git a/tests/test_imports.py b/tests/test_imports.py
index 65d915cf..3141d37e 100644
--- a/tests/test_imports.py
+++ b/tests/test_imports.py
@@ -4,6 +4,7 @@ def test_regular_imports():
     from vanna.anthropic.anthropic_chat import Anthropic_Chat
     from vanna.base.base import VannaBase
     from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
+    from vanna.hf.hf import Hf
     from vanna.local import LocalContext_OpenAI
     from vanna.marqo.marqo import Marqo_VectorStore
     from vanna.mistral.mistral import Mistral
@@ -20,6 +21,7 @@ def test_shortcut_imports():
     from vanna.anthropic import Anthropic_Chat
     from vanna.base import VannaBase
     from vanna.chromadb import ChromaDB_VectorStore
+    from vanna.hf import Hf
     from vanna.marqo import Marqo_VectorStore
     from vanna.mistral import Mistral
     from vanna.ollama import Ollama